import numpy as np
import matplotlib.pyplot as plt
import itertools
from pprint import pprint
import pickle
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.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
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("../utils/")
import circuit_utils, zne_utils, tomography_utils, sgs_algorithm
importlib.reload(circuit_utils)
importlib.reload(zne_utils)
importlib.reload(tomography_utils)
importlib.reload(sgs_algorithm)
from circuit_utils import *
from zne_utils import *
from tomography_utils import *
from sgs_algorithm import *

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

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

scale_factors = [1.0, 2.0, 3.0]

# 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>)
make_initial_state(qc, "110")  # DO NOT MODIFY (|q_5,q_3,q_1> = |110>)

subspace_encoder_init110(qc, targets=[0, 1, 2]) # encode
trotterize(qc, trot_gate, num_steps, targets=[1, 2]) # Simulate time evolution under H_heis3 Hamiltonian
subspace_decoder_init110(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), ")")

t3_zne_qcs[-3].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-community', group='ibmquantumawards', project='open-science-22')
# print("provider:", provider)
# backend = provider.get_backend("ibmq_jakarta")
print(str(backend))

shots = 1 << 13
reps = 8 # unused
jobs = []
for _ in range(reps):
    #! CHECK: run t3_zne_qcs, with optimization_level = 0 and straightforward initial_layout
    job = execute(t3_zne_qcs, backend, shots=shots, optimization_level=0)
    print('Job ID', job.job_id())
    jobs.append(job)

# QREM
qr = QuantumRegister(num_qubits, name="calq")
meas_calibs, state_labels = complete_meas_cal(qr=qr, circlabel='mcal')
# we have to feed initial_layout to calibration matrix
cal_job = execute(meas_calibs, backend=backend, shots=shots, optimization_level=3, initial_layout = initial_layout)
print('Job ID', cal_job.job_id())
meas_calibs[0].draw("mpl")

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

filename = "job_ids_" + str(backend) + "_100step_" + dt_now.strftime('%Y%m%d_%H%M%S') + "_.pkl"
print(filename)

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

# with open(filename, "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_" + str(backend) + "_" + dt_now.strftime('%Y%m%d_%H%M%S') + "_.pkl", "wb") as f:
#     pickle.dump(backend.properties(), f)

retrieved_jobs = jobs
retrieved_cal_job = cal_job

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

target_state = (One^One^Zero).to_matrix() # DO NOT CHANGE!!!

fids = []
for job in retrieved_jobs:
    # mit_results = meas_fitter.filter.apply(job.result())
    rho = StateTomographyFitter(job.result(), t3_zne_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)))