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

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

# 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

st_qcs[13].draw("mpl")

# Pauli Twirling
# TODO: 一度挙動を調べる
def pauli_twirling(circ: str) -> QuantumCircuit:
    """
    そのまま使う: 修正は後回し
    """
    #! qasm ベタ書き
    def apply_pauli(num: int, qb: int) -> str:
        if (num == 0):
            return f'id q[{qb}];\n'
        elif (num == 1):
            return f'x q[{qb}];\n'
        elif (num == 2):
            return f'y q[{qb}];\n'
        else:
            return f'z q[{qb}];\n'

    paulis = [(i,j) for i in range(0,4) for j in range(0,4)]
    paulis.remove((0,0))
    paulis_map = [(0, 1), (3, 2), (3, 3), (1, 1), (1, 0), (2, 3), (2, 2), (2, 1), (2, 0), (1, 3), (1, 2), (3, 0), (3, 1), (0, 2), (0, 3)]

    new_circ = ''
    ops = circ.qasm().splitlines(True) #! 生のqasmコードを持ってきてる: オペレータに分解
    for op in ops:
        if (op[:2] == 'cx'): # can add for cz, etc.
            num = random.randrange(len(paulis)) #! permute paulis
            qbs = re.findall('q\[(.)\]', op)
            new_circ += apply_pauli(paulis[num][0], qbs[0])
            new_circ += apply_pauli(paulis[num][1], qbs[1])
            new_circ += op
            new_circ += apply_pauli(paulis_map[num][0], qbs[0])
            new_circ += apply_pauli(paulis_map[num][1], qbs[1])
        else:
            new_circ += op
    return qiskit.circuit.QuantumCircuit.from_qasm_str(new_circ)

# ZNE
# 3種類の実行方法
# 1. state tomography回路全体をzne
# 2. state tomography前の回路をzne
# 3. 2-qubitの状態だけzne
# 今回は1を実装する
def zne_wrapper(qcs, scale_factors = [1.0, 2.0, 3.0]):
    """
    """
    folded_qcs = [] #! ZNE用の回路
    for qc in qcs:
        folded_qcs.append([mitiq.zne.scaling.fold_gates_at_random(qc, scale) for scale in scale_factors]) #! ここでmitiqを使用
    folded_qcs = list(itertools.chain(*folded_qcs)) #! folded_qcsを平坦化
    folded_qcs = [pauli_twirling(circ) for circ in folded_qcs] #! 後からPauli Twirlingを施す!
    return folded_qcs

# subspace expansion
# TODO: 実装はとりあえず後回し
def se_wrapper(qcs):
    """
    """
    folded_qcs = []
    for qc in qcs:
        pass
    return folded_qcs

from qiskit.test.mock import FakeJakarta
backend = FakeJakarta()
# backend = Aer.get_backend("qasm_simulator")
# 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)
cal_results = cal_job.result()
meas_fitter = CompleteMeasFitter(cal_results, state_labels, circlabel='mcal')

def make_basis(n):
    basis = np.zeros((1 << n, n), dtype="int")
    return basis

def measurement_basis(v):
    I = np.array([[1,0],[0,1]], dtype="complex")
    X = np.array([[0,1],[1,0]], dtype="complex")
    Y = np.array([[0,-1j],[1j,0]], dtype="complex")
    Z = np.array([[1,0],[0,-1]], dtype="complex")
    ret = I
    for i in v:
        if i == 0:
            ret = np.kron(ret, I)
        elif i == 1:
            ret = np.kron(ret, X)
        elif i == 2:
            ret = np.kron(ret, Y)
        elif i == 3:
            ret = np.kron(ret, Z)
        else:
            raise Exception
    return ret

def state_tomography(n, expvals, vs):
    rho = np.zeros((1 << n, 1 << n), dtype="complex")
    for i, v in enumerate(vs):
        rho += expvals[i] * measurement_basis(v)
    return rho / (1 << n)

# Compute the state tomography based on the st_qcs quantum circuits and the results from those ciricuits
def state_tomo(expvals, 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

stf = StateTomographyFitter(job.result(), st_qcs)
stf._data

from IPython.core.debugger import Pdb; Pdb().set_trace()
data, basis_matrix, weights = stf._fitter_data(True, 0.5)

next(iter(stf._data))

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