import numpy as np
import matplotlib.pyplot as plt
import itertools
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
from qiskit.transpiler.passes import RemoveBarriers

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

# Import mitiq for zne
import mitiq

# 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, to_instruction = True):
    qc = QuantumCircuit(2)
    qc.rx(2 * dt, 0)
    qc.rz(2 * dt, 1)
    qc.h(1)
    qc.cx(1, 0)
    qc.rz(-2 * dt, 0)
    qc.rx(-2 * dt, 1)
    qc.rz(2 * dt, 1)
    qc.cx(1, 0)
    qc.h(1)
    qc.rz(2 * dt, 0)
    return qc.to_instruction() if to_instruction else qc

trotter_gate(np.pi / 6, to_instruction=False).draw("mpl")

def make_initial_state(qc, initial_state):
    """
    logical qubit index
    little endian
    """
    for i, state in enumerate(initial_state):
        if state == "1":
            qc.x(i)

qc = QuantumCircuit(3)
make_initial_state(qc, "101")
qc.draw("mpl")

def subspace_encoder(qc, targets):
    """
    naive method, can be optimized for init state |110>
    little endian
    """
    n = qc.num_qubits
    qc.cx(targets[2],targets[1])
    qc.cx(targets[1],targets[0])

def subspace_encoder_init110(qc, targets):
    """
    optimized encoder for init state |110>
    endian: |q_0, q_1, q_2> (little endian)
    encode |110> to |0>|10>
    """
    n = qc.num_qubits
    qc.x(targets[0])

def subspace_decoder(qc, targets):
    """
    naive method
    little endian
    """
    n = qc.num_qubits
    qc.cx(targets[1], targets[0])
    qc.cx(targets[2], targets[1])

qc = QuantumCircuit(3)
make_initial_state(qc, "110")
subspace_encoder_init110(qc, targets=[0,1,2])
qc.draw("mpl")

def trotterize(qc, trot_gate, num_steps, targets):
    for _ in range(num_steps):
        qc.append(trot_gate, qargs = targets)

qc = QuantumCircuit(3)
make_initial_state(qc, "110")
subspace_encoder_init110(qc, targets=[0,1,2])
qc.barrier()
trotterize(qc, trotter_gate(np.pi / 6), 1, targets=[1, 2])
qc = transpile(qc, optimization_level = 3, basis_gates=["sx", "rz", "cx"])
qc.draw("mpl")

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

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

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

st_qcs_list = []
# Number of trotter steps
max_trotter_step = 50  ### CAN BE >= 4
trotter_steps = list(range(1, max_trotter_step + 1, 3))
for num_steps in trotter_steps:
    print("trotter step: ", num_steps)
    # Initialize quantum circuit for 3 qubits
    qr = QuantumRegister(num_qubits, name="lq")
    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>)
    make_initial_state(qc, "101")  # DO NOT MODIFY (|q_5,q_3,q_1> = |110>)
    
    subspace_encoder(qc, targets=[0, 1, 2]) # encode
    trotterize(qc, trot_gate, num_steps, targets=[1, 2]) # Simulate time evolution under H_heis3 Hamiltonian
    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})

    # Generate state tomography circuits to evaluate fidelity of simulation
    st_qcs = state_tomography_circuits(qc, [0, 1, 2][::-1])  #! state tomography requires === BIG ENDIAN === 
    
    # remove barriers
    st_qcs = [RemoveBarriers()(qc) for qc in st_qcs]
    
    # optimize circuit
    t3_st_qcs = transpile(st_qcs, optimization_level=3, basis_gates=["sx", "cx", "rz"], initial_layout=initial_layout)
    st_qcs_list.append(t3_st_qcs)

st_qcs_list[-1][-1].draw("mpl")

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

# IBMQ.load_account()
# provider = IBMQ.get_provider(hub='ibm-q-utokyo', group='internal', project='hirashi-jst')
# provider = IBMQ.get_provider(hub='ibm-q-community', group='ibmquantumawards', project='open-science-22')
# print("provider:", provider)
# backend = provider.get_backend("ibmq_jakarta")

jobs = []
shots = 1 << 13
# Number of trotter steps
for i, num_steps in enumerate(trotter_steps):
    print("trotter step: ", num_steps)
    # execute: reps = 1
    job = execute(st_qcs_list[i], backend, shots=shots)
    print('Job ID', job.job_id())
    jobs.append(job)
    print()

# QREM
qr = QuantumRegister(num_qubits)
meas_calibs, state_labels = complete_meas_cal(qr=qr, circlabel='mcal')
cal_job = execute(meas_calibs, backend=backend, shots=shots, optimization_level=3, initial_layout = initial_layout)
print('Job ID', cal_job.job_id())

cal_results = cal_job.result()
meas_fitter = CompleteMeasFitter(cal_results, state_labels, circlabel='mcal')

results = []
for job in jobs:
    results.append( job.result() )

mit_results = []
for job in jobs:
    mit_results.append( meas_fitter.filter.apply(job.result()) )

print(len(results), len(mit_results))

# 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^Zero^One).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
raw_fids = []
for result in results:
    fid = state_tomo(result, st_qcs)
    raw_fids.append(fid)
    
# print('state tomography fidelity = {:.4f} \u00B1 {:.4f}'.format(np.mean(fids), np.std(fids)))

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

plt.clf()
plt.style.use('ggplot')
plt.figure(dpi=200)
plt.title("state fidelity from Trotter step 1 to "+str(trotter_steps))
plt.plot(trotter_steps, raw_fids, label="raw fidelity")
plt.plot(trotter_steps, fids, label="fidelity after QREM")
plt.xlabel("number of trotter steps")
plt.ylabel("fidelity")
plt.grid(linestyle='dotted')

for step, fid in zip(trotter_steps, raw_fids):
    print(step, fid)

for step, fid in zip(trotter_steps, fids):
    print(step, fid)

