# This cell is added by sphinx-gallery
# It can be customized to whatever you like
%matplotlib inline

import pennylane as qml
import numpy as np
import tensorflow as tf

dev = qml.device('cirq.simulator', wires=4)

def real(angles, **kwargs):
    qml.RY(0.27740551, wires=0)
    qml.PauliX(wires =0)
    qml.CRY(0.20273270, wires=[0, 1])
    qml.PauliX(wires =0)
    qml.CRY(np.pi/2, wires=[0, 1])
    qml.CRY(np.pi/2, wires=[0, 2])
    qml.PauliX(wires =0)
    qml.CRY(1.42492, wires=[0, 2])
    qml.PauliX(wires =0)

def generator(w, **kwargs):
    qml.Hadamard(wires=0)
#     qml.Hadamard(wires=1)
    
    qml.RX(w[0], wires=0)
    qml.RX(w[1], wires=1)
    qml.RX(w[2], wires=2)
    
    qml.RY(w[3], wires=0)
    qml.RY(w[4], wires=1)
    qml.RY(w[5], wires=2)
    
    qml.RZ(w[6], wires=0)
    qml.RZ(w[7], wires=1)
    qml.RZ(w[8], wires=2)
    
    qml.CNOT(wires=[0, 1])
    qml.CNOT(wires=[1, 2])
    
    qml.RX(w[9], wires=0)
    qml.RY(w[10], wires=0)
    qml.RZ(w[11], wires=0)
    
    qml.RX(w[12], wires=1)
    qml.RY(w[13], wires=1)
    qml.RZ(w[14], wires=1)


def discriminator(w):
    qml.Hadamard(wires=1)

    qml.RX(w[0], wires=1)
    qml.RX(w[1], wires=2)
    qml.RX(w[2], wires=3)
    
    qml.RY(w[3], wires=1)
    qml.RY(w[4], wires=2)
    qml.RY(w[5], wires=3)
    
    qml.RZ(w[6], wires=1)
    qml.RZ(w[7], wires=2)
    qml.RZ(w[8], wires=3)
    
    qml.CNOT(wires=[1, 2])
    qml.CNOT(wires=[2, 3])
    
    qml.RX(w[9], wires=2)
    qml.RY(w[10], wires=2)
    qml.RZ(w[11], wires=2)
    
    qml.RX(w[12], wires=3)
    qml.RY(w[13], wires=3)
    qml.RZ(w[14], wires=3)

@qml.qnode(dev, interface="tf")
def real_disc_circuit(phi, theta, omega, disc_weights):
    real([phi, theta, omega])
    discriminator(disc_weights)
    return qml.expval(qml.PauliZ(3))


@qml.qnode(dev, interface="tf")
def gen_disc_circuit(gen_weights, disc_weights):
    generator(gen_weights)
    discriminator(disc_weights)
    return qml.expval(qml.PauliZ(3))

def prob_real_true(disc_weights):
    true_disc_output = real_disc_circuit(phi, theta, omega, disc_weights)
    # convert to probability
    prob_real_true = (true_disc_output + 1) / 2
    return prob_real_true


def prob_fake_true(gen_weights, disc_weights):
    fake_disc_output = gen_disc_circuit(gen_weights, disc_weights)
    # convert to probability
    prob_fake_true = (fake_disc_output + 1) / 2
    return prob_fake_true


def disc_cost(disc_weights):
    cost = prob_fake_true(gen_weights, disc_weights) - prob_real_true(disc_weights)
    return cost


def gen_cost(gen_weights):
    return -prob_fake_true(gen_weights, disc_weights)

phi = np.pi / 6
theta = np.pi / 2
omega = np.pi / 7
np.random.seed(0)
eps = 1e-2
init_gen_weights = np.array([np.pi] + [0] * 14) + \
                   np.random.normal(scale=eps, size=(15,))
init_disc_weights = np.random.normal(size=(15,))

gen_weights = tf.Variable(init_gen_weights)
disc_weights = tf.Variable(init_disc_weights)
print(gen_weights)

opt = tf.keras.optimizers.SGD(0.4)

cost = lambda: disc_cost(disc_weights)

for step in range(100):
    opt.minimize(cost, disc_weights)
    if step % 5 == 0:
        cost_val = cost().numpy()
        print("Step {}: cost = {}".format(step, cost_val))

print("Prob(real classified as real): ", prob_real_true(disc_weights).numpy())

print("Prob(fake classified as real): ", prob_fake_true(gen_weights, disc_weights).numpy())

cost = lambda: gen_cost(gen_weights)

for step in range(100):
    opt.minimize(cost, gen_weights)
    if step % 5 == 0:
        cost_val = cost().numpy()
        print("Step {}: cost = {}".format(step, cost_val))

print("Prob(fake classified as real): ", prob_fake_true(gen_weights, disc_weights).numpy())

print("Discriminator cost: ", disc_cost(disc_weights).numpy())

obs = [qml.PauliX(0), qml.PauliY(0), qml.PauliZ(0)]

bloch_vector_real = qml.map(real, obs, dev, interface="tf")
bloch_vector_generator = qml.map(generator, obs, dev, interface="tf")

print("Real Bloch vector: {}".format(bloch_vector_real([phi, theta, omega])))
print("Generator Bloch vector: {}".format(bloch_vector_generator(gen_weights)))

print(gen_weights)

from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute, Aer
w = gen_weights.numpy()
q = QuantumRegister(3)
c = ClassicalRegister(3)
qc = QuantumCircuit(q, c)
qc.h(q[0])
qc.rx(w[0], q[0])
qc.rx(w[1], q[1])
qc.rx(w[2], q[2])

qc.ry(w[3], q[0])
qc.ry(w[4], q[1])
qc.ry(w[5], q[2])

qc.rz(w[6], q[0])
qc.rz(w[7], q[1])
qc.rz(w[8], q[2])

qc.cx(q[0], q[1])
qc.cx(q[1], q[2])

qc.rx(w[9], q[0])
qc.ry(w[10], q[0])
qc.rz(w[11], q[0])

qc.rx(w[12], q[1])
qc.ry(w[13], q[1])
qc.rz(w[14], q[1])

job = execute(qc, backend=Aer.get_backend("statevector_simulator"))
vec = job.result().get_statevector()
for i in vec:
    print(i)

