import numpy as np
import matplotlib.pyplot as plt
plt.style.use("ggplot")
from pprint import pprint
import pickle
import time
import datetime
%matplotlib inline

# Import qubit states Zero (|0>) and One (|1>), and Pauli operators (X, Y, Z)
from qiskit.opflow import Zero, One, X, Y, Z, I
from qiskit.opflow.state_fns.dict_state_fn import DictStateFn
from qiskit.quantum_info import DensityMatrix
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, IBMQ, execute, transpile, Aer
from qiskit.tools.monitor import job_monitor
from qiskit.circuit import Parameter
from qiskit.transpiler.passes import RemoveBarriers

# Import QREM package
from qiskit.ignis.mitigation.measurement import complete_meas_cal, CompleteMeasFitter
from qiskit.ignis.mitigation import expectation_value

# Import mitiq for zne
# unused for this file
import mitiq

# Import state tomography modules
from qiskit.ignis.verification.tomography import state_tomography_circuits, StateTomographyFitter
from qiskit.quantum_info import state_fidelity

import sys
import importlib
sys.path.append("../../solutions/utils/")
import circuit_utils, zne_utils, tomography_utils
importlib.reload(circuit_utils)
importlib.reload(zne_utils) # unused for this file
importlib.reload(tomography_utils) # unused for this file
from circuit_utils import *
from zne_utils import zne_wrapper, zne_decoder
from tomography_utils import expvals_to_valid_rho

from qiskit.test.mock import FakeJakarta
# backend = FakeJakarta()
backend = Aer.get_backend("qasm_simulator")

# Returns the matrix representation of the XXX Heisenberg model for 3 spin-1/2 particles in a line
def H_heis3():
    # Interactions (I is the identity matrix; X, Y, and Z are Pauli matricies; ^ is a tensor product)
    XXs = (I^X^X) + (X^X^I)
    YYs = (I^Y^Y) + (Y^Y^I)
    ZZs = (I^Z^Z) + (Z^Z^I)
    
    # Sum interactions
    H = XXs + YYs + ZZs
    
    # Return Hamiltonian
    return H

# Returns the matrix representation of U_heis3(t) for a given time t assuming an XXX Heisenberg Hamiltonian for 3 spins-1/2 particles in a line
def U_heis3(t):
    # Compute XXX Hamiltonian for 3 spins in a line
    H = H_heis3()
    
    # Return the exponential of -i multipled by time t multipled by the 3 spin XXX Heisenberg Hamilonian 
    return (t * H).exp_i()

# Define array of time points
ts = np.linspace(0, np.pi, 100)
target_state = DictStateFn({"000": 0,
                            "001": 0,
                            "010": 1,
                            "011": 0,
                            "100": 0,
                            "101": 0,
                            "110": 1,
                            "111": 0}, 1 / np.sqrt(2))

print(DensityMatrix(target_state.to_density_matrix()).is_valid())

# Compute probability of remaining in |110> state over the array of time points
 # ~target_state gives the bra of the initial state (<110|)
 # @ is short hand for matrix multiplication
 # U_heis3(t) is the unitary time evolution at time t
 # t needs to be wrapped with float(t) to avoid a bug
 # (...).eval() returns the inner product <110|U_heis3(t)|110>
 #  np.abs(...)**2 is the modulus squared of the innner product which is the expectation value, or probability, of remaining in |110>
probs_010_110 = [np.abs((~target_state @ U_heis3(float(t)) @ target_state).eval())**2 for t in ts]

# Plot evolution of |110>
plt.figure(dpi=200)
plt.plot(ts, probs_010_110)
plt.xlabel('time')
plt.ylabel(r'probability of state $|010\rangle + |110\rangle$')
plt.title(r'Evolution of state $|010\rangle + |110\rangle$ under $H_{Heis3}$')
# plt.grid()
plt.show()

num_qubits = 3

# The final time of the state evolution
target_times = [i * np.pi / 20 for i in range(0, 20 + 1)]
pprint(target_times)

# Parameterize variable t to be evaluated at t=pi later
dt = Parameter('t')

# Convert custom quantum circuit into a gate
trot_gate = trotter_gate(dt)

# initial layout
initial_layout = [5,3,1]

# Number of trotter steps
num_steps = 100
print("trotter step: ", num_steps)

scale_factors = [1.0, 2.0, 3.0] # unused for this file

shots = 1 << 13
reps = 8 # unused

target_state = np.zeros((1 << num_qubits, 1 << num_qubits))
target_state[2, 2] = target_state[6, 6] = target_state[2, 6] = target_state[6, 2] = 0.5
target_state = DensityMatrix(target_state)
target_state.draw("latex")

prob_list = []

for target_time in target_times:
    
    print("target_time: ", target_time)
    t1 = time.perf_counter()
    
    # Initialize quantum circuit for 3 qubits
    qr = QuantumRegister(num_qubits, name="q")
    cr = ClassicalRegister(num_qubits, name="c")
    qc = QuantumCircuit(qr, cr)

    # Prepare initial state (remember we are only evolving 3 of the 7 qubits on jakarta qubits (q_5, q_3, q_1) corresponding to the state |110>)
    qc.h(0)
    qc.x(1)

    general_subspace_encoder(qc, targets=[0, 1, 2]) # encode
    trotterize(qc, trot_gate, num_steps, targets=[1, 2]) # Simulate time evolution under H_heis3 Hamiltonian
    general_subspace_decoder(qc, targets=[0, 1, 2]) # decode
    
    qc.x(1)
    qc.h(0)
    
    qc.measure(qr, cr[::-1])
    
    # Evaluate simulation at target_time (t=pi) meaning each trotter step evolves pi/trotter_steps in time
    qc = qc.bind_parameters({dt: target_time / num_steps})
    print("created qc")
    
    # optimize circuit
    t3_qc = transpile(qc, optimization_level=3, basis_gates=["sx", "cx", "rz"])
    t3_qc = transpile(t3_qc, optimization_level=3, basis_gates=["sx", "cx", "rz"])
    print("created t3_qc (depth:", t3_qc.depth(), ")")

    t3_qc = transpile(t3_qc, optimization_level=0, basis_gates=["sx", "cx", "rz"], initial_layout=initial_layout)
    print("created t3_qc (depth:", t3_qc.depth(), ")")
    
    job = execute(t3_qc, backend, shots=shots, optimization_level=0)
    print('Job ID', job.job_id())
    
    prob = job.result().get_counts().get("0" * num_qubits, 0) / shots
    prob_list.append(prob)
    
    t2 = time.perf_counter()
    print(r'probability of state |010> + |110> = {:.4f}'.format(prob))
    print("time:", t2 - t1)
    print()

plt.figure(dpi=200)
plt.plot(ts, probs_010_110, label="theoretical prediction")
plt.scatter(target_times, prob_list, c="green", zorder=4)
plt.legend(("theoretical prediction", "proposed method"))
plt.xlabel('time')
plt.ylabel(r'probability of state $|010\rangle + |110\rangle$')
plt.title(r'Evolution of state $|010\rangle + |110\rangle$ under $H_{Heis3}$')
plt.show()

import qiskit.tools.jupyter
%qiskit_version_table

fid_list = []

for target_time in target_times:
    
    print("target_time: ", target_time)
    t1 = time.perf_counter()
    
    # Initialize quantum circuit for 3 qubits
    qr = QuantumRegister(num_qubits, name="q")
    qc = QuantumCircuit(qr)

    # Prepare initial state (remember we are only evolving 3 of the 7 qubits on jakarta qubits (q_5, q_3, q_1) corresponding to the state |110>)
    qc.h(0)
    qc.x(1)

    general_subspace_encoder(qc, targets=[0, 1, 2]) # encode
    trotterize(qc, trot_gate, num_steps, targets=[1, 2]) # Simulate time evolution under H_heis3 Hamiltonian
    general_subspace_decoder(qc, targets=[0, 1, 2]) # decode

    # Evaluate simulation at target_time (t=pi) meaning each trotter step evolves pi/trotter_steps in time
    qc = qc.bind_parameters({dt: target_time / num_steps})
    print("created qc")

    # Generate state tomography circuits to evaluate fidelity of simulation
    st_qcs = state_tomography_circuits(qc, [0, 1, 2][::-1])  #! state tomography requires === BIG ENDIAN === 
    print("created st_qcs (length:", len(st_qcs), ")")

    # remove barriers
    st_qcs = [RemoveBarriers()(qc) for qc in st_qcs]
    print("removed barriers from st_qcs")

    # optimize circuit
    t3_st_qcs = transpile(st_qcs, optimization_level=3, basis_gates=["sx", "cx", "rz"])
    t3_st_qcs = transpile(t3_st_qcs, optimization_level=3, basis_gates=["sx", "cx", "rz"])
    print("created t3_st_qcs (length:", len(t3_st_qcs), ")")

    # zne wrapping
    zne_qcs = t3_st_qcs
    # zne_qcs = zne_wrapper(t3_st_qcs, scale_factors = scale_factors, pt = False)
    # print("created zne_qcs (length:", len(zne_qcs), ")")

    # optimization_level must be 0
    # feed initial_layout here to see the picture of the circuits before casting the job
    t3_zne_qcs = transpile(zne_qcs, optimization_level=0, basis_gates=["sx", "cx", "rz"], initial_layout=initial_layout)
    print("created t3_zne_qcs (length:", len(t3_zne_qcs), ")")
    
    job = execute(t3_zne_qcs, backend, shots=shots, optimization_level=0)
    print('Job ID', job.job_id())
    
    rho = StateTomographyFitter(job.result(), t3_zne_qcs).fit(method='lstsq')
    fid = state_fidelity(rho, target_state)
    fid_list.append(fid)
    
    t2 = time.perf_counter()
    print('state tomography fidelity = {:.4f}'.format(fid))
    print("time:", t2 - t1)
    print()