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

import numpy as np
from qiskit import QuantumCircuit
from qiskit.opflow import I, X, Y, Z
import matplotlib.pyplot as plt

from volta.observables import sample_hamiltonian
from volta.hamiltonians import BCS_hamiltonian

EPSILONS = [3, 3, 3, 3]
V = -2
hamiltonian = 1 / 2 * (Z ^ I) + 1 / 2 * (Z ^ Z) #BCS_hamiltonian(EPSILONS, V)
print(hamiltonian)

eigenvalues, _ = np.linalg.eigh(hamiltonian.to_matrix())
print(f"Eigenvalues: {eigenvalues}")

def create_circuit(n: int) -> list:
    return [QuantumCircuit(n) for _ in range(2)]

n = hamiltonian.num_qubits
init_states = create_circuit(n)

def copy_unitary(list_states: list) -> list:
    out_states = []
    for state in list_states:
        out_states.append(state.copy())
    return out_states

import textwrap

def apply_initialization(list_states: list) -> None:
    for ind, state in enumerate(list_states):
        b = bin(ind)[2:]
        if len(b) != n:
            b = '0'*(n - len(b)) + b
        spl = textwrap.wrap(b, 1)
        for qubit, val in enumerate(spl):
            if val == '1':
                state.x(qubit)

apply_initialization(init_states)

initialization = copy_unitary(init_states)

from qiskit.circuit.library import TwoLocal
ansatz = TwoLocal(hamiltonian.num_qubits, ['ry','rz'], 'cx', reps=2)

def apply_ansatz(ansatz: QuantumCircuit, list_states: list) -> None:
    for states in list_states:
        states.append(ansatz, range(n))

apply_ansatz(ansatz, init_states)

init_states[-1].draw('mpl')

def _apply_varform_params(ansatz, params: list):
        """Get an hardware-efficient ansatz for n_qubits
        given parameters.
        """
        # Define variational Form
        var_form = ansatz

        # Get Parameters from the variational form
        var_form_params = sorted(var_form.parameters, key=lambda p: p.name)

        # Check if the number of parameters is compatible
        assert len(var_form_params) == len(params), "The number of parameters don't match"

        # Create a dictionary with the parameters and values
        param_dict = dict(zip(var_form_params, params))

        # Assing those values for the ansatz
        wave_function = var_form.assign_parameters(param_dict)

        return wave_function

from qiskit import BasicAer
from qiskit.utils import QuantumInstance

def cost_function(params:list) -> float:
        backend = BasicAer.get_backend('qasm_simulator')
        backend = QuantumInstance(backend, shots=10000)
        cost = 0
        
        w = np.arange(len(init_states), 0, -1)
        
        for i, state in enumerate(init_states):
            qc = _apply_varform_params(state, params)

            # Hamiltonian
            hamiltonian_eval = sample_hamiltonian(hamiltonian=hamiltonian, 
                                                 ansatz=qc, 
                                                 backend=backend)
            cost += w[i] * hamiltonian_eval

        return cost

from qiskit.aqua.components.optimizers import COBYLA
optimizer = COBYLA()

n_parameters = len(init_states[0].parameters)

params = np.random.rand(n_parameters)

optimal_params, mean_energy, n_iters = optimizer.optimize(num_vars=n_parameters, 
                                                        objective_function=cost_function, 
                                                        initial_point=params)

mean_energy, sum(eigenvalues[:2])

# Optimized first ansatz
ansatz_1 = _apply_varform_params(ansatz, optimal_params)
ansatz_1.name = 'U(θ)'

def cost_function_ind(ind: int, params:list) -> float:
        backend = BasicAer.get_backend('qasm_simulator')
        backend = QuantumInstance(backend, shots=10000)
        cost = 0
        
        # Define Ansatz
        qc = _apply_varform_params(init_states[ind], params)

        # Hamiltonian
        hamiltonian_eval = sample_hamiltonian(hamiltonian=hamiltonian, 
                                             ansatz=qc, 
                                             backend=backend)
        cost += hamiltonian_eval

        
        return cost

energies = []
for i in range(len(init_states)):
    energies.append(cost_function_ind(ind=i, params=optimal_params))
energies = np.array(energies)

for i in range(len(init_states)):
    print(f"Value for the {i} state: {energies[i]}\t expected energy {eigenvalues[i]}")

