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

import numpy as np
import matplotlib.pyplot as plt

# Qiskit
from qiskit.aqua.operators import I, X, Y, Z
from qiskit import BasicAer
from qiskit.aqua.components.optimizers import COBYLA

# VOLTA
from VOLTA.VQD import VQD
from VOLTA.utils import classical_solver
from VOLTA.Hamiltonians import BCS_hamiltonian

%load_ext autoreload
%autoreload 2

def SPT_hamiltonian(g:float):
    gzz = 2*(1 - g**2)
    gx = (1 + g)**2
    gzxz = (g - 1)**2
    return - (gzz*(Z^Z^I^I^I) + gzz*(I^Z^Z^I^I) + gzz*(I^I^Z^Z^I) + gzz*(I^I^I^Z^Z)) - (gx*(X^I^I^I^I) + gx*(I^X^I^I^I) + gx*(I^I^X^I^I) + gx*(I^I^I^X^I) + gx*(I^I^I^I^X)) + (gzxz*(Z^X^Z^I^I) + gzxz*(I^Z^X^Z^I) + gzxz*(I^I^Z^X^Z))

g = 1
hamiltonian = SPT_hamiltonian(g)
print(hamiltonian)

eigenvalues, _ = classical_solver(hamiltonian)
print(f"Energy Density: {eigenvalues[0]/3}")
print(f"Expected Energy: {-2*(1 + g**2)}")

energy = []
n_points = 10
x = np.linspace(-1, 1, n_points)
for g in x:
    hamiltonian = SPT_hamiltonian(g)
    eigenvalues, _ = classical_solver(hamiltonian)
    energy.append(eigenvalues[0])

plt.plot(x, np.array(energy)/5, label='energy')
ideal = lambda g:-2*(1 + g**2)
plt.plot(x, ideal(x), label='expected')
plt.legend()
plt.show()

import numpy as np
from VOLTA.Observables import sample_hamiltonian
from tqdm import tqdm
optimizer = COBYLA()
backend = BasicAer.get_backend('qasm_simulator')

def simulation(n_points:int,
              backend,
              optimizer) -> (np.array, np.array):
        g = np.linspace(-1, 1, n_points)
        result = []
        for g_val in tqdm(g):
            hamiltonian = SPT_hamiltonian(g_val)
            Algo = VQD(hamiltonian=hamiltonian,
                        n_excited_states=0,
                        beta=0.,
                        optimizer=optimizer,
                        backend=backend)
            Algo.run(0)      
            vqd_energies = Algo.energies
            vqd_states = Algo.states
            
            string_order = Z^Y^X^Y^Z
            res = sample_hamiltonian(hamiltonian=string_order,
                   backend=backend,
                   ansatz=vqd_states[0])
            result.append(res)
            
        return g, np.array(result)

g, result = simulation(n_points=10,
                       backend=backend,
                       optimizer=optimizer)
plt.plot(g, result, 'o');

from qiskit.aqua import QuantumInstance
from qiskit.circuit.library import TwoLocal

optimizer = COBYLA(maxiter=10000)
backend = QuantumInstance(backend=BasicAer.get_backend('qasm_simulator'),
                              shots=10000)

# Define ansatz
ansatz = TwoLocal(hamiltonian.num_qubits, ['ry','rz'], 'cx', reps=2)

# Run algorithm
Algo = VQD(hamiltonian=hamiltonian,
            ansatz=ansatz,
            n_excited_states=0,
            beta=0.,
            optimizer=optimizer,
            backend=backend)

Algo.run(0)      
vqd_energies = Algo.energies
vqd_states = Algo.states
print(f'Energy: {vqd_energies[0]}')

from VOLTA.Observables import sample_hamiltonian
#string_order = Z^Y^X^Y^Z
string_order = I^I^X^I^I
sample_hamiltonian(hamiltonian=string_order,
                   backend=backend,
                   ansatz=vqd_states[0])

4*np.abs(g)/(1 + np.abs(g))**2

