import time
import numpy as np
import qiskit
from qiskit.opflow import X,Z,I
from qiskit.opflow.state_fns import StateFn, CircuitStateFn

from modules.utils import *

import matplotlib.pyplot as plt

from scipy import sparse
import scipy.sparse.linalg.eigen.arpack as arp
%matplotlib inline

# first need a gate that is the exponential of ZZ, which shouldnt be too difficult since it can be done analytically

L = 2

qreg = qiskit.QuantumRegister(L, 'q')
creg = qiskit.ClassicalRegister(2, 'c')
circ = qiskit.QuantumCircuit(qreg, creg)

circ.rz()

tau = qiskit.circuit.Parameter("tau")
diaglist = [tau, -tau, -tau, tau]
diaglist

dt = 0.05

diaglist = [np.exp(dt), np.exp(-dt), np.exp(-dt), np.exp(dt)]

qiskit.circuit.library.Diagonal(diaglist)





import numpy as np
from qiskit.opflow import I, X, Y, Z, H, CX, Zero, ListOp, PauliExpectation, PauliTrotterEvolution, CircuitSampler, MatrixEvolution, Suzuki
from qiskit.circuit import Parameter
from qiskit import BasicAer

L=5
anti = -1
gx = 1e-2
gz = 1e-2



two_qubit_H2 = QHIsing(L,anti,np.float32(gx),np.float32(gz))

two_qubit_H2 = 1 * (X^X^X^Z^X + Z^Z^Z^Z^X)
two_qubit_H2

evo_time = Parameter('θ')
evolution_op = (evo_time * two_qubit_H2).exp_i()

h2_measurement = StateFn(two_qubit_H2).adjoint()
#print(h2_measurement)

evo_and_meas = h2_measurement @ evolution_op @ Zero
#print(evo_and_meas)

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

bound = trotterized_op.bind_parameters({evo_time: .5})

diagonalized_meas_op = PauliExpectation().convert(trotterized_op)

evo_time_points = [1,2,3,4]
h2_trotter_expectations = diagonalized_meas_op.bind_parameters({evo_time: evo_time_points})

h2_trotter_expectations.eval()







# Ops
H = QHIsing(L,anti,np.float32(gx),np.float32(gz))
H_meas = StateFn(H).adjoint()
mag = QMag(L,anti)

qiskit.opflow.One^5

qiskit.opflow.VectorStateFn(qiskit.opflow.One^5)

asd = H_meas @ qiskit.opflow.One^5

asd.eval()

evo_time = qiskit.circuit.Parameter('θ') # note that we will have to bind it laterevo_time = Parameter('θ')

evolution_op = (evo_time  * H).exp_i() #* 1j
# expH = suzi.convert(H)

evo_and_meas = H_meas @ evolution_op

trotterized_op = qiskit.opflow.evolutions.PauliTrotterEvolution(trotter_mode=qiskit.opflow.evolutions.Suzuki(order=2, reps=1)).convert(evo_and_meas)

bound = trotterized_op.bind_parameters({evo_time: .5})

bound.eval()

diagonalized_meas_op = qiskit.opflow.PauliExpectation().convert(trotterized_op)

h2_trotter_expectations = diagonalized_meas_op.bind_parameters({evo_time: 1})

h2_trotter_expectations.eval()

H.__class__, evolution_op.__class__

def imaginary_time_evolution(tau, L, gx, gz, anti=-1):
    H = QHIsing(L,anti,np.float32(gx),np.float32(gz))
    evo_time = qiskit.circuit.Parameter('tau') # note that we will have to bind it later
    evolution_op = (tau * 1j * H).exp_i()
    trotterized_op = qiskit.opflow.evolutions.PauliTrotterEvolution(trotter_mode=qiskit.opflow.evolutions.Suzuki(order=2, reps=1)).convert(H_meas @ evolution_op)
    return trotterized_op

trott_op = imaginary_time_evolution(tau=1, L=5, gx=1e-2, gz=1e-2, anti=-1)

trott_op.bind_param({evo_})

bound = trott_op.bind_parameters({evo_time: 1})

bound.to_circuit().draw()

mag_res = ~StateFn(mag) @ StateFn(trott_op)

mag_res.eval()

bound.eval()

bound = trotterized_op.bind_parameters({evo_time: tau})

# this is a dummy circuit for the sake of checking how I can combine the evolved state with a circuit later
circ = qiskit.circuit.library.TwoLocal(5,"rx","cz", reps=1,entanglement="linear")

circ.draw()

circ_fn = CircuitStateFn(circ)

print(circ_fn)

evolved_processed = evolution_op @ circ_fn # first the imaginary time evolution, then the dummy processing
print(evolved_processed )

trotterized_op = qiskit.opflow.evolutions.PauliTrotterEvolution(trotter_mode=qiskit.opflow.evolutions.Suzuki(order=2, reps=1)).convert(evolved_processed)
print(trotterized_op)

bound = trotterized_op.bind_parameters({evo_time: .5j})

bound.to_circuit().draw()

# ED result
Smag = Mag(L,anti)
init_state, E, ham = ising_groundstate(L, anti, np.float(gx), np.float(gz))
Sen = E
Smags = init_state.T.conj()@Smag@init_state #Magnetization with Numpy results

