import numpy as np
import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 16})  # enlarge matplotlib fonts

# Import qubit states Zero (|0>) and One (|1>), and Pauli operators (X, Y, Z)
from qiskit.opflow import Zero, One, I, X, Y, Z
from qiskit import QuantumCircuit, QuantumRegister, IBMQ, execute, transpile, Aer
from qiskit.providers.aer import QasmSimulator
from qiskit.tools.monitor import job_monitor
from qiskit.circuit import Parameter

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

# Suppress warnings
import warnings
warnings.filterwarnings('ignore')

def trotter_gate(dt):
    qc = QuantumCircuit(2)
    qc.h(1)
    qc.cx(1,0)
    qc.h(1)
    qc.rz(- 2 * dt, 1)
    qc.rz(dt, 0)
    qc.h(1)
    qc.cx(1,0)
    qc.h(1)
    qc.rx(dt, [1])

    qc.rz(-dt, [0,1])
    qc.rx(-dt, [0,1])
    qc = qc.reverse_bits()
    return qc.to_instruction()

qc = QuantumCircuit(2)
qc.h(1)
qc.cx(1,0)
qc.h(1)
qc.rz(- 2 * np.pi / 6, 1)
qc.rz(np.pi / 6, 0)
qc.h(1)
qc.cx(1,0)
qc.h(1)
qc.rx(np.pi / 6, [1])
qc.barrier()
qc.rz(-np.pi / 6, [0,1])
qc.rx(-np.pi / 6, [0,1])
qc = qc.reverse_bits()
qc.draw('mpl')

# Combine subcircuits into a single multiqubit gate representing a single trotter step
num_qubits = 3

# 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)

# YOUR TROTTERIZATION GOES HERE -- FINISH (end of example)

# The final time of the state evolution
target_time = np.pi

# Number of trotter steps
trotter_steps = 12  ### CAN BE >= 4

# Initialize quantum circuit for 3 qubits
qr = QuantumRegister(7)
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.x([3])  # DO NOT MODIFY (|q_5,q_3,q_1> = |110>)

# Simulate time evolution under H_heis3 Hamiltonian
for _ in range(trotter_steps):
    qc.append(Trot_gate, [qr[3], qr[5]])

qc.cx(qr[3], qr[1])
qc.cx(qr[5], qr[3])
    
# Evaluate simulation at target_time (t=pi) meaning each trotter step evolves pi/trotter_steps in time
qc = qc.bind_parameters({dt: target_time/trotter_steps})

shots = 8192
reps = 8

# WE USE A NOISELESS SIMULATION HERE
backend = Aer.get_backend('qasm_simulator')

# Generate state tomography circuits to evaluate fidelity of simulation
st_qcs = state_tomography_circuits(qc, [qr[1], qr[3], qr[5]])

# Display circuit for confirmation
# st_qcs[-1].decompose().draw()  # view decomposition of trotter gates
st_qcs[-1].draw("mpl")  # only view trotter gates

jobs = []
for _ in range(reps):
    # execute
    job = execute(st_qcs, backend, shots=shots, optimization_level=0)
    print('Job ID', job.job_id())
    jobs.append(job)

# Compute the state tomography based on the st_qcs quantum circuits and the results from those ciricuits
def state_tomo(result, st_qcs):
    # The expected final state; necessary to determine state tomography fidelity
    target_state = (One^One^Zero).to_matrix()  # DO NOT MODIFY (|q_5,q_3,q_1> = |110>)
    # Fit state tomography results
    tomo_fitter = StateTomographyFitter(result, st_qcs)
    rho_fit = tomo_fitter.fit(method='lstsq')
    # Compute fidelity
    fid = state_fidelity(rho_fit, target_state)
    return fid

# Compute tomography fidelities for each repetition
fids = []
for job in jobs:
    fid = state_tomo(job.result(), st_qcs)
    fids.append(fid)
    
print('state tomography fidelity = {:.4f} \u00B1 {:.4f}'.format(np.mean(fids), np.std(fids)))

from qiskit.compiler import transpile

t0_qc = transpile(qc, optimization_level=0, basis_gates=["sx","rz","cx"])
t0_qc.draw("mpl")

t1_qc = transpile(qc, optimization_level=1, basis_gates=["sx","rz","cx"])
t1_qc.draw("mpl")

t2_qc = transpile(qc, optimization_level=2, basis_gates=["sx","rz","cx"])
t2_qc.draw("mpl")

t3_qc = transpile(qc, optimization_level=3, basis_gates=["sx","rz","cx"])
t3_qc.draw("mpl")

st_qcs = state_tomography_circuits(t2_qc, [qr[1], qr[3], qr[5]])
jobs = []
for _ in range(reps):
    # execute
    job = execute(st_qcs, backend, shots=shots)
    print('Job ID', job.job_id())
    jobs.append(job)

fids = []
for job in jobs:
    fid = state_tomo(job.result(), st_qcs)
    fids.append(fid)
    
print('state tomography fidelity = {:.4f} \u00B1 {:.4f}'.format(np.mean(fids), np.std(fids)))

st_qcs = state_tomography_circuits(t3_qc, [qr[1], qr[3], qr[5]])
jobs = []
for _ in range(reps):
    # execute
    job = execute(st_qcs, backend, shots=shots)
    print('Job ID', job.job_id())
    jobs.append(job)

fids = []
for job in jobs:
    fid = state_tomo(job.result(), st_qcs)
    fids.append(fid)
    
print('state tomography fidelity = {:.4f} \u00B1 {:.4f}'.format(np.mean(fids), np.std(fids)))

