from qiskit.utils import algorithm_globals
from qiskit import QuantumCircuit, Aer, execute
from qiskit.circuit.library import RealAmplitudes
from qiskit.primitives import Sampler
from qiskit_machine_learning.connectors import TorchConnector
from qiskit_machine_learning.neural_networks import SamplerQNN
from torch import nn
from torch.optim import Adam
import torch, time, argparse, sys
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from datetime import datetime
from tqdm import tqdm

matplotlib.rcParams['mathtext.fontset'] = 'stix'
matplotlib.rcParams['font.family'] = 'STIXGeneral'
width=0.75
color='black'
fontsize=28
ticksize=22
figsize=(10,10)

# For this example, I have chosen a power-law with power -7/3
def targ_prob(xs):
    px = xs**(-7./3)
    return px/np.sum(px)

# Set the number of qubits
n = 4
# Number of dimensions (currently hard coded to 1 but can be expanded)
num_dim = 1
# Number of evaluation points
num_discrete_values = 2**n

# Range of x and the intervals between points
xmin, xmax = 40., 168.
dx = (xmax-xmin)/num_discrete_values

# Full set of x values
xs = np.arange(xmin, xmax, dx)
# Evaluate the target distribution over this range
prob_data = targ_prob(xs)

coords = np.linspace(-2, 2, num_discrete_values)
grid_elements = np.expand_dims(coords,axis=1)

# Set the number of samples S we wish to take from the PQC and target distribution 
shots=10000

# Set the random seed used in the simulation (the outcome is highly sensitive to this)
rseed = 1696004649
# Uncomment line below to change the random seed for the training
#rseed = int(datetime.now().timestamp())
algorithm_globals.random_seed = rseed
_ = torch.manual_seed(rseed) 

# Define the sampler object
sampler = Sampler(options={"shots": shots, "seed": algorithm_globals.random_seed})

qc = QuantumCircuit(n)
# Begin by initiallising the qubits in an equal superposition
qc.h(qc.qubits);

# Set the number of layers L
reps = 12

ansatz = RealAmplitudes(n, reps=reps)
qc.compose(ansatz, inplace=True)

def create_generator() -> TorchConnector:
    qnn = SamplerQNN(
        circuit=qc,
        sampler=sampler,
        input_params=[],
        weight_params=qc.parameters,
        sparse=False,
    )

    initial_weights = algorithm_globals.random.random(qc.num_parameters)
    return TorchConnector(qnn, initial_weights) 

class Discriminator(nn.Module):
    def __init__(self, input_size):
        super(Discriminator, self).__init__()

        self.linear_input = nn.Linear(input_size, 20)
        self.leaky_relu = nn.LeakyReLU(0.2)
        self.linear20 = nn.Linear(20, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        x = self.linear_input(input)
        x = self.leaky_relu(x)
        x = self.linear20(x)
        x = self.sigmoid(x)
        return x

generator = create_generator()
discriminator = Discriminator(num_dim)

def adversarial_loss(input, target, w):
    bce_loss = target * torch.log(input) + (1 - target) * torch.log(1 - input)
    weighted_loss = w * bce_loss
    total_loss = -torch.sum(weighted_loss)
    return total_loss

lr=0.01
b1=0.7
b2=0.999

generator_optimizer = Adam(generator.parameters(), lr=lr, betas=(b1, b2), weight_decay=0.005)
discriminator_optimizer = Adam(discriminator.parameters(), lr=lr, betas=(b1, b2), weight_decay=0.005)

generator_loss_values = []
discriminator_loss_values = []
entropy_values = []

# Set the number of epochs
n_epochs=1000

num_qnn_outputs = num_discrete_values**num_dim

# Resume training if already performed some iterations
n_epochs = n_epochs-len(entropy_values)
        
start = time.time()
for epoch in tqdm(range(n_epochs)):

    # Define torch target objects
    valid = torch.ones(num_qnn_outputs, 1, dtype=torch.float)
    fake = torch.zeros(num_qnn_outputs, 1, dtype=torch.float)

    # Define torch tensor representing the target probability distribution
    real_dist = torch.tensor(prob_data, dtype=torch.float).reshape(-1, 1)

    # Define torch tensor representing the grid over which samples are taken
    samples = torch.tensor(grid_elements, dtype=torch.float)
    disc_value = discriminator(samples)

    # Generate data
    gen_dist = generator(torch.tensor([])).reshape(-1, 1)

    # Train generator
    generator_optimizer.zero_grad()
    generator_loss = adversarial_loss(disc_value, valid, gen_dist)

    # store for plotting
    generator_loss_values.append(generator_loss.detach().item())

    generator_loss.backward(retain_graph=True)
    generator_optimizer.step()

    # Train Discriminator
    discriminator_optimizer.zero_grad()

    real_loss = adversarial_loss(disc_value, valid, real_dist)
    fake_loss = adversarial_loss(disc_value, fake, gen_dist.detach())
    discriminator_loss = (real_loss + fake_loss) / 2

    # Store for plotting
    discriminator_loss_values.append(discriminator_loss.detach().item())

    discriminator_loss.backward()
    discriminator_optimizer.step()

    # Retreive trained parameters
    with torch.no_grad():
        generated_weights = generator.weight.detach().numpy().reshape((reps+1,n))
    
    # Build circuit to execute given trained parameters
    circ = QuantumCircuit(n)
    circ.h(circ.qubits)

    # Implement parameterised quantum circuit equvalent to RealAmplitude
    for rep in np.arange(reps+1):
        for i in np.arange(n):
            circ.ry(generated_weights[rep,i],circ.qubits[i])
        for i in np.arange(n):
            if i!=n-1 and rep!=reps:
                circ.cx(circ.qubits[i], circ.qubits[i+1])
    
    # Produce resulting state vector
    backend = Aer.get_backend('statevector_simulator')
    job = execute(circ, backend)
    result = job.result()
    state_vector = np.asarray(result.get_statevector())
    
    # Calculate value equivalent to mismatch from fidelity
    entropy_value = 1-np.sqrt(np.abs(np.dot(np.sqrt(prob_data),np.conjugate(state_vector)))**2)
    entropy_values.append(entropy_value)
            
print(np.round(100.*(epoch/n_epochs),2),'%','Mismatch:',"{:.2e}".format(entropy_value),'Generator loss:', np.round(generator_loss_values[-1],2), 'Discriminator loss:', np.round(discriminator_loss_values[-1],2))

elapsed = time.time() - start
print(f"Fit in {elapsed:0.2f} sec")

fig = plt.figure(figsize=np.array(figsize)*[1,2])
ax = fig.add_subplot(211)

plt.plot(generator_loss_values, lw=2, color='black')
plt.plot(discriminator_loss_values, lw=2, color='grey')
ax.set_ylabel('Loss function', fontsize=fontsize);
ax.tick_params(axis='both', labelsize=ticksize);

ax = fig.add_subplot(212)

plt.plot(np.log10(entropy_values), lw=2, color='black')
ax.set_ylabel(r'log$_{10}(1-\sqrt{\mathcal{F}})$', fontsize=fontsize);
ax.set_xlabel('Iterations', fontsize=fontsize);
ax.tick_params(axis='both', labelsize=ticksize);

with torch.no_grad():
    generated_weights = generator.weight.detach().numpy().reshape((reps+1,n))

circ = QuantumCircuit(n)
circ.h(circ.qubits)

for rep in np.arange(reps+1):
    for i in np.arange(n):
        circ.ry(generated_weights[rep,i],circ.qubits[i])
    for i in np.arange(n):
        if i!=n-1 and rep!=reps:
            circ.cx(circ.qubits[i], circ.qubits[i+1])

backend = Aer.get_backend('statevector_simulator')
job = execute(circ, backend)
result = job.result()
state_vector = np.asarray(result.get_statevector())

fidelity = np.abs(np.dot(np.sqrt(prob_data),np.conjugate(state_vector)))**2

mismatch = 1. - np.sqrt(fidelity)

print('Fidelity:', fidelity, 'Mismatch:', mismatch)

target_state = np.sqrt(targ_prob(xs))

fig, ax = plt.subplots(2, 1, figsize=figsize, gridspec_kw={'height_ratios': [1.5, 1]})

color1 = 'black'
color2 = 'grey'

ax[0].scatter(xs, state_vector.real, color=color1, lw=2)
ax[0].plot(xs, state_vector.real, color=color1, ls='-', lw=1)
ax[0].plot(xs, target_state, color=color1, ls='--', lw=2)

ax[0].tick_params(axis='both', labelsize=ticksize)
ax[1].set_xlabel(r'$x$', fontsize=fontsize)
ax[0].set_ylabel(r'Amplitude', fontsize=fontsize)
ax[0].set_xlim(xmin-1,xmax+1)

ax[0].set_xticks([])

ax[1].scatter(xs, state_vector.real-target_state.real, color=color1, lw=2)
ax[1].set_ylabel(r'$\Delta$ Amplitude', fontsize=fontsize)
ax[1].tick_params(axis='both', labelsize=ticksize)
ax[1].set_xlim(xmin-1,xmax+1)

fig.tight_layout()

plt.show()