# General imports
import numpy as np
import matplotlib.pyplot as plt

# Qiskit imports
import qiskit as qk
from qiskit.utils import QuantumInstance

# Qiskit Machine Learning imports
import qiskit_machine_learning as qkml
from qiskit_machine_learning.neural_networks import CircuitQNN
from qiskit_machine_learning.connectors import TorchConnector

# PyTorch imports
import torch
from torch import Tensor

# OpenAI Gym import
import gym

# Custom Deep Q-Learning code import
from dqn_definitions import *

# Fix seed for reproducibility
seed = 42
np.random.seed(seed)
torch.manual_seed(seed);

# To get smooth animations on Jupyter Notebooks. 
# Note: these plotting function are taken from https://github.com/ageron/handson-ml2
import matplotlib as mpl
import matplotlib.animation as animation
mpl.rc('animation', html='jshtml')

def update_scene(num, frames, patch):
    patch.set_data(frames[num])
    return patch,

def plot_animation(frames, repeat=False, interval=40):
    fig = plt.figure()
    patch = plt.imshow(frames[0])
    plt.axis('off')
    anim = animation.FuncAnimation(
        fig, update_scene, fargs=(frames, patch),
        frames=len(frames), repeat=repeat, interval=interval)
    plt.close()
    return anim

# Select the number of qubits (one qubit per state variable of the environment)
num_qubits = 4

# Generate the Parametrized Quantum Circuit (note the reuploading and repetitions flags)
qc = parametrized_circuit(num_qubits = num_qubits, 
                          reuploading = True,
                          reps = 6)

# Fetch the Parameters from the circuit and divide them in Inputs (X) and Trainable Parameters (params)
# The first four parameters are for the inputs
X = list(qc.parameters)[: num_qubits]

# The remaining ones are the trainable weights of the quantum neural network
params = list(qc.parameters)[num_qubits:]

qc.draw()

# Define the quantum instance to be used for simulating the quantum circuit
qi = QuantumInstance(qk.Aer.get_backend('statevector_simulator'))

# Create the QNN from the parametrized quantum circuit defined above
qnn = CircuitQNN(qc, input_params=X, weight_params=params, 
                 quantum_instance = qi)

# Interface the QNN with PyTorch
initial_weights = (2*np.random.rand(qnn.num_weights) - 1) # <---Initialization weights at random
quantum_nn = TorchConnector(qnn, initial_weights)

env = gym.make("CartPole-v1")
input_shape = [4] # == env.observation_space.shape
n_outputs = 2 # == env.action_space.n

# Classical trainable preprocessing (encoding) 
encoding = encoding_layer()

# Classical trainable postprocessing
exp_val = exp_val_layer()

# Stack the classical and quantum layers together 
model = torch.nn.Sequential(encoding, 
                            quantum_nn, 
                            exp_val)

# Load pre-trained weights (check if the file exists):
try:
    model.load_state_dict(torch.load("./model_best_weights_6reps_longtrain.pth"))
except: 
    print("No pre-trained weights found. Looks like you have to train from scratch...")

# Print the optimal weights of the agent 
model.state_dict()

# Set the environment seed for reproducibility 
env.seed(42)

# 200 is the target score for considering the environment solved
max_reward = 200

# Run the quantum agent
frames = []
state = env.reset()
for step in range(max_reward):
    # Select action
    with torch.no_grad():
        Q_values = model(Tensor(state)).numpy()
    action = np.argmax(Q_values)
    
    # Perform the action
    state, reward, done, info = env.step(action)
    if done:
        print("End at step:", step)
        break
    
    print(f"Step {step}, Q-values {Q_values}, Action {action}", end = "\r")
    
    img = env.render(mode="rgb_array")
    frames.append(img)

plot_animation(frames)

# Choose the number of shots used to estimate the expectation values in the QNN circuit
n_shots = 1024

# Define the quantum instance for the qasm_simulator
qi_qasm = QuantumInstance(qk.Aer.get_backend('qasm_simulator'), shots = n_shots)

noisy_qnn = CircuitQNN(qc, input_params=X, weight_params=params, 
                 quantum_instance = qi_qasm)

initial_weights = (2*np.random.rand(qnn.num_weights) - 1)
noisy_quantum_nn = TorchConnector(noisy_qnn, initial_weights)

# Exactly as before...
encoding = encoding_layer()
exp_val = exp_val_layer()

noisy_model = torch.nn.Sequential(encoding, 
                            noisy_quantum_nn, 
                            exp_val)

# Load pre-trained weights:
try:
    noisy_model.load_state_dict(torch.load("./model_best_weights_6reps_longtrain.pth"))
except: 
    print("No pre-trained weights found. Looks like you have to train from scratch...")

noisy_model.state_dict()

env.seed(42)

number_of_episodes = 5 

print(f"Number of shots = {n_shots}")
print("-----------------------")

final_step = []
for i in range(number_of_episodes):
    state = env.reset()
    for step in range(200):
        print(f"\r Episode = {i}, Step = {step}", end = "")
        with torch.no_grad():
            Q_values = noisy_model(Tensor(state[np.newaxis])).numpy()
        action = np.argmax(Q_values[0])

        state, reward, done, info = env.step(action)
        # print(action, state, model(Tensor([state])))
        if done:
            print(", End at step:", step)
            break

    final_step.append(step+1)

final_steps = np.array(final_step)
mean_score = np.mean(final_steps)
std_score = np.std(final_steps)
print(f"Mean reward = {mean_score} ± {std_score}")

num_shots = [1024, 2048, 4096, 8192]
mean_reward = [115.1, 191, 195, 199.7]
std_reward = [49.9, 27.2, 10.9, 0.3]

import seaborn as sns
sns.set_theme()

cmap = plt.get_cmap('tab20c')

fig = plt.figure(figsize=(8,5))
plt.axhline([200], ls = 'dashed', c=cmap(9))
plt.text(6000,205, s='Max reward', c=cmap(8))

plt.axhline([195], ls = 'dashed', c=cmap(13))
plt.text(6000,175, s='Min reward for\nsolving environment', c=cmap(12))

plt.errorbar(num_shots, mean_reward, yerr=std_reward, fmt='o', c = cmap(4), capsize=5, ecolor = cmap(5))
plt.ylim([0,230])
plt.xticks([1024, 2048, 4096, 8192])
plt.xlabel("Number of shots")
plt.ylabel("Final reward")

plt.tight_layout()
plt.show()