%matplotlib inline
import sys
sys.path.append("..")
import compress
from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit, execute, Aer
import numpy as np
from functools import partial
from qiskit.aqua.components.optimizers import SPSA
from multideterminant_prep import PrepareMultiDeterminantState as pmds
import seaborn as sns
import matplotlib.pyplot as plt

get_ipython().run_line_magic('pylab', 'inline')
pylab.rcParams['figure.figsize'] = (12, 6)
pylab.rcParams.update({'font.size': 10})

in_strings = ["101", "011", "110"]
in_weights = [4, 7, 2]
circuit = pmds(in_weights, in_strings)
circuit.draw(scale=0.5, style={'fold':100}, output='mpl')

# Configuration
n_shots  = 1000
n_trials = 300
n_layers = 3
n_params = circuit.width() * n_layers * 3

# Build variable bounds
variable_bounds_single = (0., 2*np.pi)
variable_bounds = [variable_bounds_single] * n_params
initial_point = np.random.uniform(low=variable_bounds_single[0],
                                  high=variable_bounds_single[1],
                                  size=(n_params,)).tolist()
circuit.draw(scale=0.5, style={'fold':100}, output='mpl')

print("Circuit has a depth of: {}".format(circuit.depth()))

final_circuit = compress.swap_test_with_compression_model(circuit, initial_point)
final_circuit.draw(scale=0.5, style={'fold':100}, output='mpl')

results = compress.cross_validate_qnn_depth(circuit, n_shots, n_trials, n_layers, run=0)
results

sns.lineplot(x='iteration', y='fidelity', data=results.fidelity.to_dataframe().reset_index())
plt.title("Fidelity improvement with optimiser iterations")

