import numpy as np
from qiskit.aqua.operators import I, X, Y, Z
import matplotlib.pyplot as plt
%load_ext autoreload
%autoreload 2

def BCS_hamiltonian(epsilons, V):
    return ((1/2*epsilons[0]*I^Z) + (1/2*epsilons[1]*Z^I) + V*(1/2*(X^X) + 1/2*(Y^Y)))

EPSILONS = [1,1]
V = 2
hamiltonian = BCS_hamiltonian(EPSILONS, V)
print(hamiltonian)

from qiskit.aqua.operators import PauliTrotterEvolution, CircuitSampler, MatrixEvolution, Suzuki
from qiskit.aqua.operators import StateFn, Plus, Zero, One, Minus, PauliExpectation, CX
from qiskit.circuit import Parameter
from qiskit import BasicAer

# Initial State 
init_state = (Plus^Zero)

# Time evolution operator
evolution_param = Parameter("t")
evolution_op = (evolution_param * hamiltonian).exp_i()

# Measurement of the Hamiltonian with the time evolved state
evo_and_meas = StateFn(hamiltonian).adjoint() @ evolution_op @ init_state

# Trotterization
trotterized_op = PauliTrotterEvolution(trotter_mode=Suzuki(order=2, reps=1)).convert(evo_and_meas)

# Diagonalize the measurement
diagonalized_meas_op = PauliExpectation().convert(trotterized_op)

# Create an array to represent time and apply to the circuit
n_points = 10
time_array = list(np.linspace(0, 1/V, n_points))
expectations = diagonalized_meas_op.bind_parameters({evolution_param: time_array})

# Define Backend
backend = BasicAer.get_backend("statevector_simulator")

# Use CircuitSampler for getting energy values
sampler = CircuitSampler(backend=backend)
sampled_expectations = sampler.convert(expectations)
sampled_energies = sampled_expectations.eval()
print(f"Time:\n {time_array}")
print(f"Sampled energies after Trotterization:\n {np.real(sampled_energies)}")

plt.plot(time_array, np.real(sampled_energies))
plt.show()

