# Let's import all functions
import numpy as np
import scipy as sp
from qiskit import Aer
from qiskit.opflow import PauliTrotterEvolution, StateFn, PauliExpectation
from qiskit.opflow import CircuitSampler, PauliOp
from qiskit.opflow import I, X, Y, Z, Zero, One, Plus, Minus
from qiskit.circuit import Parameter

hamiltonian = (Z^Z)

evo_time = Parameter('t')
evolution_op = (evo_time*hamiltonian).exp_i()
print(evolution_op) 

num_time_slices = 1
trotterized_op = PauliTrotterEvolution(
                    trotter_mode='trotter', 
                    reps=num_time_slices).convert(evolution_op)
trotterized_op.to_circuit().draw()

hamiltonian = (X^X)
evo_time = Parameter('t')
evolution_op = (evo_time*hamiltonian).exp_i()
num_time_slices = 1
trotterized_op = PauliTrotterEvolution(
                    trotter_mode='trotter', 
                    reps=num_time_slices).convert(evolution_op)
trotterized_op.to_circuit().draw()

hamiltonian = (Y^Y)
evo_time = Parameter('t')
evolution_op = (evo_time*hamiltonian).exp_i()
num_time_slices = 1
trotterized_op = PauliTrotterEvolution(
                    trotter_mode='trotter', 
                    reps=num_time_slices).convert(evolution_op)
trotterized_op.to_circuit().draw()

hamiltonian = (Z^Z^Z^Z)
evo_time = Parameter('t')
evolution_op = (evo_time*hamiltonian).exp_i()
num_time_slices = 1
trotterized_op = PauliTrotterEvolution(
                    trotter_mode='trotter', 
                    reps=num_time_slices).convert(evolution_op)
trotterized_op.to_circuit().draw()

hamiltonian = (Z^Z) + (X^X)
observable = (X^X)
evo_time = Parameter('t')
evolution_op = (evo_time*hamiltonian).exp_i()

observable_measurement = StateFn(observable).adjoint()

eigenvalues, eigenstates = np.linalg.eigh(hamiltonian.to_matrix())
initial_state = StateFn(eigenstates[0])
print(initial_state.to_circuit_op())

evo_and_measure = observable_measurement @ evolution_op @ initial_state
print(evo_and_measure)

num_time_slices = 1
trotterized_op = PauliTrotterEvolution(
                    trotter_mode='trotter', 
                    reps=num_time_slices).convert(evo_and_measure)
print(trotterized_op)

# Let's calculate expectation values
diagonalized_meas_op = PauliExpectation().convert(trotterized_op)
print(diagonalized_meas_op)

evo_time_points = [0.5, 0.75]
hamiltonian_trotter_expectations = diagonalized_meas_op.bind_parameters({evo_time: evo_time_points})
print(f"Observable at time {evo_time_points}: {np.round(hamiltonian_trotter_expectations.eval(), 3)}")

sampler = CircuitSampler(backend=Aer.get_backend("qasm_simulator"))
# sampler.quantum_instance.run_config.shots = 1000
sampled_trotter_exp_op = sampler.convert(hamiltonian_trotter_expectations)
sampled_trotter_energies = sampled_trotter_exp_op.eval()
print(f"Energies: {np.round(np.real(sampled_trotter_energies),3)}")

import matplotlib.pyplot as plt

def run_hs(shots):
    sampler = CircuitSampler(backend=Aer.get_backend("qasm_simulator"))
    sampler.quantum_instance.run_config.shots = shots
    hamiltonian_trotter_expectations = diagonalized_meas_op.bind_parameters({evo_time: 0.5})
    sampled_trotter_exp_op = sampler.convert(hamiltonian_trotter_expectations)
    sampled_trotter_energies = sampled_trotter_exp_op.eval()
    return np.real(sampled_trotter_energies)

n_shots = [100, 1000, 2000, 5000, 7000, 10000]
exp = []
for shots in n_shots:
    exp.append(run_hs(shots))

plt.plot(n_shots, exp, 'o', label=f"t={0.5}")
plt.hlines(y=0., xmin=min(n_shots), xmax=max(n_shots) + 1, colors='red')
plt.xlabel("# shots", size=16)
plt.ylabel(r"$ \langle O \rangle $", size=16)
plt.show()