import sys
sys.path.append('../../')

import numpy as np
from numpy import linalg as la
import matplotlib.pyplot as plt
from qiskit import QuantumCircuit

from compilers.aqc_demo.cnot_structures import full_conn, spin
from compilers.aqc_demo.parametric_circuit import ParametricCircuit
from compilers.aqc_demo.test_sample_data import ORIGINAL_CIRCUIT, INITIAL_THETAS
from compilers.aqc_demo.aqc_optimizers import FISTAOptimizer, GDOptimizer
from compilers.aqc_demo.aqc import AQC

num_qubits = int(np.log2(ORIGINAL_CIRCUIT.shape[0]))
num_qubits

def cnots_to_qc(num_qubits: int, cnots) -> QuantumCircuit:
    circuit = QuantumCircuit(num_qubits)
    for ii, elements in enumerate(cnots[0]):
        circuit.cx(elements - 1, cnots[1][ii] - 1)

    return circuit

cnots = spin(num_qubits, full_conn(num_qubits))
cnots_qc = cnots_to_qc(num_qubits, cnots)
cnots_qc.draw(output="mpl")

cnots_qc.depth()

target_circuit = ParametricCircuit(num_qubits, cnots)

initial_thetas_qc = target_circuit.make_circuit(INITIAL_THETAS, reverse=True)
initial_thetas_qc.draw(output="mpl")

initial_thetas_qc.depth()

def approximation_error(expected, actual) -> float:
    return .5 * (la.norm(expected - actual, 'fro') ** 2)

target_unitary = target_circuit.compute_unitary(INITIAL_THETAS)
error = approximation_error(ORIGINAL_CIRCUIT, target_unitary)
error

optimizer = FISTAOptimizer(method="nesterov", maxiter=5000, eta=0.1, tol=0.01, eps=0.0, reg=0.7, group=True)
thetas, objective, gradient, _ = optimizer.optimize(ORIGINAL_CIRCUIT, INITIAL_THETAS, target_circuit)

target_unitary = target_circuit.compute_unitary(thetas)
error = approximation_error(ORIGINAL_CIRCUIT, target_unitary)
error

compressed_cnots, compressed_thetas, spar = target_circuit.compress(thetas, synth=False)
compressed_target_circuit = ParametricCircuit(num_qubits, compressed_cnots)
compressed_qc_circuit = compressed_target_circuit.make_circuit(compressed_thetas, reverse=True)
compressed_qc_circuit.draw(output="mpl")

compressed_qc_circuit.depth()

optimizer = GDOptimizer(method="nesterov", maxiter=5000, eta=0.1, tol=0.01, eps=0.0)
thetas, objective, gradient, thetas_min = optimizer.optimize(ORIGINAL_CIRCUIT, 
                                                             compressed_thetas, 
                                                             compressed_target_circuit)

num_iterations = np.shape(objective)[0]
plt.loglog(range(num_iterations), objective)
plt.xlabel('Iteration')
plt.ylabel('Approximation error')
plt.show()

target_unitary = compressed_target_circuit.compute_unitary(thetas)
error = approximation_error(ORIGINAL_CIRCUIT, target_unitary)
error

aqc = AQC(method="nesterov", maxiter=5000, eta=0.1, tol=0.01, eps=0.0, reg=0.7, group=True)
parametric_circuit, thetas = aqc.compile_unitary(ORIGINAL_CIRCUIT, cnots, INITIAL_THETAS)

target_unitary = parametric_circuit.compute_unitary(thetas)
error = approximation_error(ORIGINAL_CIRCUIT, target_unitary)
error

