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

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

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

# Number of trotter steps
trotter_steps = 15  ### 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([5])  # 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})
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"])

# Generate state tomography circuits to evaluate fidelity of simulation
st_qcs = state_tomography_circuits(t3_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

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

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

shots = 8192
reps = 8

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

# 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=[5,3,1])
print('Job ID', cal_job.job_id())

dt_now = datetime.datetime.now()
print(dt_now)

with open("jobs_jakarta_15step_" + dt_now.strftime('%Y%m%d_%H%M%S') + "_.pkl", "wb") as f:
    pickle.dump({"jobs": jobs, "cal_job": cal_job}, f)

with open("job_ids_jakarta_15step_" + dt_now.strftime('%Y%m%d_%H%M%S') + "_.pkl", "wb") as f:
    pickle.dump({"job_ids": [job.job_id() for job in jobs], "cal_job_id": cal_job.job_id()}, f)

with open("properties_jakarta" + dt_now.strftime('%Y%m%d_%H%M%S') + "_.pkl", "wb") as f:
    pickle.dump(backend.properties(), f)

filename = "job_ids_jakarta_15step_20220412_031437_.pkl" # change here
with open(filename, "rb") as f:
    job_ids_dict = pickle.load(f)
    job_ids = job_ids_dict["job_ids"]
    cal_job_id = job_ids_dict["cal_job_id"]

retrieved_jobs = []
for job_id in job_ids:
    retrieved_jobs.append(backend.retrieve_job(job_id))
retrieved_cal_job = backend.retrieve_job(cal_job_id)

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

# set the target state
target_state = (One^One^Zero).to_matrix() # DO NOT CHANGE!!!

fids = []
for job in retrieved_jobs:
    # mit_results = meas_fitter.filter.apply(job.result()) # apply QREM
    rho = StateTomographyFitter(job.result(), st_qcs).fit(method='lstsq')
    fid = state_fidelity(rho, target_state)
    fids.append(fid)
    
print('state tomography fidelity = {:.4f} \u00B1 {:.4f}'.format(np.mean(fids), np.std(fids)))

import qiskit.tools.jupyter
%qiskit_version_table

