# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.3.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

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

#
#
# Quantum Generative Adversarial Networks with Cirq + TensorFlow
# ==============================================================
#
# .. meta::
#     :property="og:description": This demo constructs and trains a Quantum
#         Generative Adversarial Network (QGAN) using PennyLane, Cirq, and TensorFlow.
#     :property="og:image": https://pennylane.ai/qml/_images/qgan3.png
#
# This demo constructs a Quantum Generative Adversarial Network (QGAN)
# (`Lloyd and Weedbrook
# (2018) <https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.121.040502>`__,
# `Dallaire-Demers and Killoran
# (2018) <https://journals.aps.org/pra/abstract/10.1103/PhysRevA.98.012324>`__)
# using two subcircuits, a *generator* and a *discriminator*. The
# generator attempts to generate synthetic quantum data to match a pattern
# of "real" data, while the discriminator tries to discern real data from
# fake data (see image below). The gradient of the discriminator’s output provides a
# training signal for the generator to improve its fake generated data.
#
# |
#
# .. figure:: ../demonstrations/QGAN/qgan.png
#     :align: center
#     :width: 75%
#     :target: javascript:void(0)
#
# |
#

# Using Cirq + TensorFlow
# ~~~~~~~~~~~~~~~~~~~~~~~
# PennyLane allows us to mix and match quantum devices and classical machine
# learning software. For this demo, we will link together
# Google's `Cirq <https://cirq.readthedocs.io/en/stable/>`_ and `TensorFlow <https://www.tensorflow.org/>`_ libraries.
#
# We begin by importing PennyLane, NumPy, and TensorFlow.
#
#

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

# We also declare a 3-qubit simulator device running in Cirq.
#
#

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


# Generator and Discriminator
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# In classical GANs, the starting point is to draw samples either from
# some “real data” distribution, or from the generator, and feed them to
# the discriminator. In this QGAN example, we will use a quantum circuit
# to generate the real data.
#
# For this simple example, our real data will be a qubit that has been
# rotated (from the starting state $\left|0\right\rangle$) to some
# arbitrary, but fixed, state.
#
#

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)


# For the generator and discriminator, we will choose the same basic
# circuit structure, but acting on different wires.
#
# Both the real data circuit and the generator will output on wire 0,
# which will be connected as an input to the discriminator. Wire 1 is
# provided as a workspace for the generator, while the discriminator’s
# output will be on wire 2.
#
#

# +
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)


# -

# We create two QNodes. One where the real data source is wired up to the
# discriminator, and one where the generator is connected to the
# discriminator. In order to pass TensorFlow Variables into the quantum
# circuits, we specify the ``"tf"`` interface.
#
#

# +
@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))


# -

# QGAN cost functions
# ~~~~~~~~~~~~~~~~~~~
#
# There are two cost functions of interest, corresponding to the two
# stages of QGAN training. These cost functions are built from two pieces:
# the first piece is the probability that the discriminator correctly
# classifies real data as real. The second piece is the probability that the
# discriminator classifies fake data (i.e., a state prepared by the
# generator) as real.
#
# The discriminator is trained to maximize the probability of
# correctly classifying real data, while minimizing the probability of
# mistakenly classifying fake data.
#
# The generator is trained to maximize the probability that the
# discriminator accepts fake data as real.
#
#

# +
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)


# -

# Training the QGAN
# ~~~~~~~~~~~~~~~~~
#
# We initialize the fixed angles of the “real data” circuit, as well as
# the initial parameters for both generator and discriminator. These are
# chosen so that the generator initially prepares a state on wire 0 that
# is very close to the $\left| 1 \right\rangle$ state.
#
#

# +
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)
# -

# We begin by creating the optimizer:
#
#

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

# In the first stage of training, we optimize the discriminator while
# keeping the generator parameters fixed.
#
#

# +
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))
# -

# At the discriminator’s optimum, the probability for the discriminator to
# correctly classify the real data should be close to one.
#
#

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

# For comparison, we check how the discriminator classifies the
# generator’s (still unoptimized) fake data:
#
#

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

# In the adversarial game we now have to train the generator to better
# fool the discriminator. For this demo, we only perform one stage of the
# game. For more complex models, we would continue training the models in an
# alternating fashion until we reach the optimum point of the two-player
# adversarial game.
#
#

# +
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))
# -

# At the optimum of the generator, the probability for the discriminator
# to be fooled should be close to 1.
#
#

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

# At the joint optimum the discriminator cost will be close to zero,
# indicating that the discriminator assigns equal probability to both real and
# generated data.
#
#

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

# The generator has successfully learned how to simulate the real data
# enough to fool the discriminator.
#
# Let's conclude by comparing the states of the real data circuit and the generator. We expect
# the generator to have learned to be in a state that is very close to the one prepared in the
# real data circuit. An easy way to access the state of the first qubit is through its
# `Bloch sphere <https://en.wikipedia.org/wiki/Bloch_sphere>`__ representation:
#
#

# +
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)


