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, ClassicalRegister, IBMQ, execute, transpile, Aer
from qiskit.providers.aer import QasmSimulator
from qiskit.tools.monitor import job_monitor
from qiskit.circuit import Parameter
from qiskit.compiler import transpile

# 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])
    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.draw('mpl')

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

# 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 = 4  ### CAN BE >= 4

# Initialize quantum circuit for 3 qubits
qr = QuantumRegister(2)
cr = ClassicalRegister(2)
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>)
# init state |10> (= |110>)
qc.x(1)  # 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[0], qr[1]])

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

qc.measure(qr, cr)

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

t0_qc = t0_qc.reverse_bits()
# t0_qc.draw("mpl")

shots = 8192
reps = 1

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

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

counts_01 = []
counts_10 = []
for trotter_steps in range(1, 15, 1):
    
    print("number of trotter steps: ", trotter_steps)
    # Initialize quantum circuit for 3 qubits
    qr = QuantumRegister(2)
    cr = ClassicalRegister(2)
    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>)
    # init state |10> (= |110>)
    qc.x(1)  # 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[0], qr[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 / trotter_steps})
    qc.measure(qr, cr)
    t0_qc = transpile(qc, optimization_level=0, basis_gates=["sx","rz","cx"])
    t0_qc = t0_qc.reverse_bits()
    print("circuit depth: ", t0_qc.depth())
    job = execute(t0_qc, backend=backend, shots=shots, optimization_level=0)
    print("pribability distribution: ", job.result().get_counts())
    counts_01.append(job.result().get_counts().get("01", 0))
    counts_10.append(job.result().get_counts().get("10", 0))
    print()

plt.plot(range(1,15), counts_10)
plt.xlabel("trotter steps")
plt.ylabel("shot counts of 10")
plt.title("counts of |10>")

plt.plot(range(1,15), counts_01)
plt.xlabel("trotter steps")
plt.ylabel("shot counts of 01")
plt.title("counts of |01>")

