!pip install pennylane

# example of training a gan on mnist
from numpy import expand_dims
from numpy import zeros
from numpy import ones
from numpy import vstack
from numpy.random import randn
from numpy.random import randint
import tensorflow as tf
from keras.datasets.mnist import load_data
from keras.optimizers import Adam
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import Reshape
from keras.layers import Flatten
from keras.layers import Conv2D
from keras.layers import Conv2DTranspose
from keras.layers import LeakyReLU
from keras.layers import Dropout
from matplotlib import pyplot
import pennylane as qml
import numpy as np
num_wires=5
dev = qml.device('default.qubit', wires=5)
# variables
phi = [np.pi] * 12
for i in range(len(phi)):
		phi[i] = phi[i] / np.random.randint(2, 12)
num_epochs = 30
eps = 1 

initial_generator_weights = np.array([np.pi] + [0] * 44) + np.random.normal(scale=eps, size=(45,))
initial_discriminator_weights = np.random.normal(scale=eps, size=(35,))

generator_weights = tf.Variable(initial_generator_weights)
discriminator_weights = tf.Variable(initial_discriminator_weights)
opt = tf.keras.optimizers.Adam(0.4)

def real_data(phi, **kwargs):
		# phi is a list with length 12
		qml.Hadamard(wires=0)
		qml.CNOT(wires=[0,1])
		qml.RX(phi[0], wires=0)
		qml.RY(phi[1], wires=0)
		qml.RZ(phi[2], wires=0)
		qml.RX(phi[3], wires=0)
		qml.RY(phi[4], wires=0)
		qml.RZ(phi[5], wires=0)
		qml.CNOT(wires=[0, 1])
		qml.RX(phi[6], wires=1)
		qml.RY(phi[7], wires=1)
		qml.RZ(phi[8], wires=1)
		qml.RX(phi[9], wires=1)
		qml.RY(phi[10], wires=1)
		qml.RZ(phi[11], wires=1)
# the discriminator acts on wires 0, 1, and 4

def discriminator_layer(w, **kwargs):
		qml.RX(w[0], wires=0)
		qml.RX(w[1], wires=1)
		qml.RX(w[2], wires=4)
		qml.RZ(w[3], wires=0)
		qml.RZ(w[4], wires=1)
		qml.RZ(w[5], wires=4)
		qml.MultiRZ(w[6], wires=[0, 1])
		qml.MultiRZ(w[7], wires=[1, 4])

def discriminator(w, **kwargs):
		qml.Hadamard(wires=0)
		qml.CNOT(wires=[0,1])
		discriminator_layer(w[:8])
		discriminator_layer(w[8:16]) 
		discriminator_layer(w[16:32])
		qml.RX(w[32], wires=4)
		qml.RY(w[33], wires=4)
		qml.RZ(w[34], wires=4)


  
def generator_layer(w):
		qml.RX(w[0], wires=0)
		qml.RX(w[1], wires=1)
		qml.RX(w[2], wires=2)
		qml.RX(w[3], wires=3)
		qml.RZ(w[4], wires=0)
		qml.RZ(w[5], wires=1)
		qml.RZ(w[6], wires=2)
		qml.RZ(w[7], wires=3)
		qml.MultiRZ(w[8], wires=[0, 1])
		qml.MultiRZ(w[9], wires=[2, 3])
		qml.MultiRZ(w[10], wires=[1, 2])
  
def generator(w, **kwargs):
		qml.Hadamard(wires=0)
		qml.CNOT(wires=[0,1])
		generator_layer(w[:11])
		generator_layer(w[11:22])
		generator_layer(w[22:33]) # includes w[22], doesnt include w[33]
		qml.RX(w[33], wires=0)
		qml.RY(w[34], wires=0)
		qml.RZ(w[35], wires=0)
		qml.RX(w[36], wires=1)
		qml.RY(w[37], wires=1)
		qml.RZ(w[38], wires=1)
		qml.CNOT(wires=[0, 1])
		qml.RX(w[39], wires=0)
		qml.RY(w[40], wires=0)
		qml.RZ(w[41], wires=0)
		qml.RX(w[42], wires=1)
		qml.RY(w[43], wires=1)
		qml.RZ(w[44], wires=1)
  
@qml.qnode(dev, interface='tf')
def real_discriminator(phi, discriminator_weights):
		real_data(phi)
		discriminator(discriminator_weights)
		return qml.expval(qml.PauliZ(4))

@qml.qnode(dev, interface='tf')
def generator_discriminator(generator_weights, discriminator_weights):
		generator(generator_weights)
		discriminator(discriminator_weights)
		return qml.expval(qml.PauliZ(4))
  
def probability_real_real(discriminator_weights):
		# probability of guessing real data as real
		discriminator_output = real_discriminator(phi, discriminator_weights) # the output of the discriminator classifying the data as real
		probability_real_real = (discriminator_output + 1) / 2
		return probability_real_real

def probability_fake_real(generator_weights, discriminator_weights):
		# probability of guessing real fake as real
		# incorrect classification
		discriminator_output = generator_discriminator(generator_weights, discriminator_weights)
		probability_fake_real = (discriminator_output + 1) / 2
		return probability_fake_real

def discriminator_cost(discriminator_weights):
		accuracy = probability_real_real(discriminator_weights) - probability_fake_real(generator_weights, discriminator_weights)
		# accuracy = correct classification - incorrect classification
		cost = -accuracy    
		return cost

def generator_cost(generator_weights):
		accuracy = probability_fake_real(generator_weights, discriminator_weights)
		# accuracy = probability that the generator fools the discriminator
		cost = -accuracy
		return cost



def train_discriminator():
		for epoch in range(num_epochs):
				cost = lambda: discriminator_cost(discriminator_weights) 
				# you need lambda because discriminator weights is a tensorflow object
				opt.minimize(cost, discriminator_weights)
				if epoch % 5 == 0:
						cost_val = discriminator_cost(discriminator_weights).numpy()
						print('Epoch  {}/{}, Cost: {}, Probability class real as real: {}'.format(epoch, num_epochs, cost_val, probability_real_real(discriminator_weights).numpy()))
				if epoch == num_epochs - 1:
						print('\n')

def train_generator():
		for epoch in range(num_epochs):
				cost = lambda: generator_cost(generator_weights)
				opt.minimize(cost, generator_weights)
				if epoch % 5 == 0:
						cost_val = generator_cost(generator_weights).numpy()
						print('Epoch  {}/{}, Cost: {}, Probability class fake as real: {}'.format(epoch, num_epochs, cost_val, probability_fake_real(generator_weights, discriminator_weights).numpy()))
				if epoch == num_epochs - 1:
						print('\n')

train_discriminator()
train_generator()



