import logging
from copy import deepcopy
import warnings
import numpy as np
from numpy import pi
import matplotlib.pyplot as plt
from qiskit.circuit import Gate
from qiskit import transpile
from qiskit.transpiler import InstructionProperties
from qiskit_utilities import (
    save_job_data,
    setup_logger,
    retrieve_expdata,
    get_cr_schedule,
    iy_drag_calibration,
    get_default_cr_params,
    read_calibration_data,
    create_direct_cnot_schedule,
    rough_rz0_correction_calibration,
    fine_rz0_correction_calibration,
    initilize_qiskit_dynamics_backend
)
session = None  # qiskit runtime session

# Setup a logging system.
setup_logger(filename=None, level=logging.INFO, stdout=True)
logger = logging.getLogger("qiskit_utilities")

# Create a qiskit-dynamics backend with two Transmon, see fake_backend.py for detail.
# fidelity
f0 = 5.188339840468758e9
f1 = 5.079647268847252e9
# anharmonicity
a0 = -0.341931669003125e9
a1 = -0.343248404293835e9
backend = initilize_qiskit_dynamics_backend(f0, f1, a0, a1, J=0.0030e9)

# Operating qubits
QUBIT_C, QUBIT_T = (0, 1)
# Multiple the amplitude by a ratio factor.
ratio = 1
# Compute the sampling time steps for the CR tomography experiments
if backend.name == 'DynamicsBackend':
    amp = 2 * pi * 0.06
    cr_times = np.arange(16*16, 1440*3, 32*4, dtype=int)
    duration = 720*2 # Used for plotting only, not tomography.
else:
    amp = get_default_cr_params(backend, QUBIT_C, QUBIT_T)[0]
    duration = backend.defaults().instruction_schedule_map.get("cx", (QUBIT_C, QUBIT_T)).duration/16 * 4 / ratio
    cr_times = 16 * np.arange(16, duration + 16, duration//30, dtype=int)

# Custom shape
gate_name = "CR-recursive-tr10-direct"
params = {
    "order": "3",  # order m for the initial pulse shape
    "t_r": 10,
    "drag_type" : "exact",  # Givens rotation, use "sw" for perturbation
    "duration": duration,  # only for plotting
    "amp": amp,
    "angle": 0.0,
    ## If real device, use the following
    # "Delta": (backend.properties().frequency(QUBIT_C) - backend.properties().frequency(QUBIT_T)) * 1.0e-9 * 2 * pi,
    # "a1": 2 * pi * backend.properties().qubit_property(QUBIT_C)["anharmonicity"][0] * 1.0e-9,
    ## If Dynamicsbackend, use this
    "Delta": 2 * pi * (f0 - f1) * 1.0e-9,
    "a1": 2 * pi * a0 * 1.0e-9,
    "drag_scale": [1., 1., 1.],
}
cr_params = params.copy()
ix_params = params.copy()
ix_params["amp"] = 0.0
ix_params["order"] = "2"
ix_params["drag_type"] = "01"  # single-derivative DRAG

initial_calibration_data = {
    "cr_params": cr_params,
    "ix_params": ix_params,
    "frequency_offset": 0.,
}

get_cr_schedule((QUBIT_C, QUBIT_T), backend, cr_params=cr_params, ix_params=ix_params).draw()

from qiskit_utilities import iterative_cr_pulse_calibration
iterative_cr_pulse_calibration(
    (QUBIT_C, QUBIT_T),
    backend,
    initial_calibration_data=initial_calibration_data,
    cr_times=cr_times,
    session=session,
    gate_name=gate_name,
    verbose=True,
    restart=True,
    max_repeat=3,
    rerun_last_calibration=True,
    shots=2048,
    mode="CR",
    )

iy_drag_calibration((QUBIT_C, QUBIT_T), backend, gate_name, cr_times, session, verbose=True, delta_beta=10.)

# Reduce cr_times by a half
if backend.name == 'DynamicsBackend':
    cr_times = np.arange(16*16, 1440*4, 32*4, dtype=int)
else:
    duration = backend.defaults().instruction_schedule_map.get("cx", (QUBIT_C, QUBIT_T)).duration/16 * 4 / ratio / 2
    cr_times = 16 * np.arange(16, duration + 16, duration//30, dtype=int)
from qiskit_utilities import iterative_cr_pulse_calibration
iterative_cr_pulse_calibration(
    (QUBIT_C, QUBIT_T),
    backend,
    cr_times,
    session,
    gate_name=gate_name,
    initial_calibration_data=None,
    verbose=True,
    threshold_MHz=0.015,
    restart=True,
    rerun_last_calibration=False,
    max_repeat=3,
    shots=8096,
    mode="IX-pi"
    )

rough_rz0_correction_calibration(backend, (QUBIT_C, QUBIT_T), gate_name, session)
fine_rz0_correction_calibration(backend, (QUBIT_C, QUBIT_T), gate_name, session, num_repeat_cnot=6)

from qiskit_experiments.library import ProcessTomography
calibration_data = read_calibration_data(backend, gate_name, (QUBIT_C, QUBIT_T))
custom_cnot_sched = create_direct_cnot_schedule(backend, (QUBIT_C, QUBIT_T), calibration_data, reverse_direction=False)
# Show the custom CNOT schedule.
custom_cnot_sched.draw()

custom_cnot = Gate("custom_cnot", 2, [])
backend_tmp = deepcopy(backend)
backend_tmp.target.add_instruction(
    custom_cnot,
    {(QUBIT_C, QUBIT_T): InstructionProperties(calibration=custom_cnot_sched)},
    name="custom_cnot",
)
qptexp = ProcessTomography(custom_cnot, physical_qubits=[QUBIT_C, QUBIT_T])

try:
    basis_gates = backend.configuration().basis_gates
except AttributeError:
    basis_gates = ["measure", "sx", "x", "rz"]
transpiled_circ_list = transpile(
    qptexp._transpiled_circuits(),
    backend=backend_tmp,
    basis_gates = basis_gates + ["custom_cnot"],
    optimization_level= 1,  # Turn off optimization to keep the custom pulse
    )

shots=8096
if session is not None:
    qpt_job = session.run(
        "circuit-runner",
        inputs={
            "circuits": transpiled_circ_list,
            "skip_transpilation": True,
            "shots": shots,
        },
    )
else:
    qpt_job = backend.run(transpiled_circ_list, shots=shots)

parameters = {
    "backend": backend.name,
    "qubits": (QUBIT_C, QUBIT_T),
    "shots": shots,
    }
logger.info(
    "Quantum process tomography job: "+
    str(qpt_job.job_id()) + "\n"
    + "\n".join([f"{key}: {val}" for key, val in parameters.items()])
    + "\n"
    )

save_job_data(qpt_job, backend=backend, parameters=parameters)

import qutip
qptdata = retrieve_expdata(qpt_job, backend=backend, exp=qptexp)

state_result = qptdata.analysis_results("state")
choi_matrix = state_result.value.data
choi_qobj = qutip.Qobj(
    choi_matrix, dims=[[[2, 2], [2, 2]], [[2, 2], [2, 2]]], superrep="choi"
)
choi_qobj = choi_qobj.permute([[1, 0], [3, 2]])  # Permute the two qubit
qutip.to_super(choi_qobj)
op_basis = [[qutip.qeye(2), qutip.sigmax(), qutip.sigmay(), qutip.sigmaz()]] * 2
chi = qutip.qpt(qutip.to_super(choi_qobj), op_basis)
qutip.qpt_plot_combined(
    chi, lbls_list=[["i", "x", "y", "z"]] * 2, title="Process Matrix"
)

from qiskit.circuit.library import CXGate
from qiskit_experiments.library import InterleavedRB


mode = "with_correction"
direct = True
if mode not in ("with_correction", "without_correction", "default"):
    raise ValueError("Unknown mode.")
backend_tmp = deepcopy(backend)
rb_gate = CXGate()
calibration_data = read_calibration_data(backend, gate_name, (QUBIT_C, QUBIT_T))
if mode != "default":
    modified_calibration_data = deepcopy(calibration_data)
    if mode == "without_correction":
        modified_calibration_data["ix_params"]["beta"] = 0.
        modified_calibration_data["frequency_offset"] = 0.
        modified_calibration_data["x_gate_frequency_offset"] = 0.
    # Direct gate
    custom_cnot_sched = create_direct_cnot_schedule(backend, (QUBIT_C, QUBIT_T), modified_calibration_data, with_separate_x=False)
    backend_tmp.target["cx"][(QUBIT_C, QUBIT_T)] = InstructionProperties(calibration=custom_cnot_sched)

    custom_cnot_sched = create_direct_cnot_schedule(backend, (QUBIT_C, QUBIT_T), modified_calibration_data, with_separate_x=False, reverse_direction=True)
    backend_tmp.target["cx"][(QUBIT_T, QUBIT_C)] = InstructionProperties(calibration=custom_cnot_sched)


seed = np.random.randint(2**16)
num_sampling = 30
lengths = [1, 10, 20, 30, 40, 50, 60]

irb_exp = InterleavedRB(
    interleaved_element=rb_gate, physical_qubits=[QUBIT_C, QUBIT_T], lengths=lengths, num_samples=num_sampling, backend=backend_tmp, seed=seed)

transpiled_circ_list = transpile(
    # One must use _transpiled_circuits because it will include the calibration
    irb_exp._transpiled_circuits(),
    backend=backend_tmp,
    basis_gates= ["x", "sx", "rz", "cx"], #+ ["custom_cnot"],
    # initial_layout = [QUBIT_C, QUBIT_T],  # Map the circuit to the target qubits
    optimization_level= 1,  # Turn off optimization to keep the custom pulse
    )
shots = 4096
rb_job = backend_tmp.run(
    transpiled_circ_list,
    shots=shots,
    )
parameters = {
    "backend": backend.name,
    "mode": mode,
    "with_separate_x": False,
    "direct": direct,
    "qubits": [QUBIT_C, QUBIT_T],
    "gate": rb_gate.name,
    "shots": shots,
    "lengths": lengths,
    "num_sampling": num_sampling,
    "seed": seed,
    "calibration_data": calibration_data,
    }
logger.info(
    "Randomized benchmarking job: "+
    str(rb_job.job_id()) + "\n"
    + "\n".join([f"{key}: {val}" for key, val in parameters.items()])
    + "\n"
    )
save_job_data(rb_job, backend=backend, parameters=parameters)


from qiskit_utilities import load_job_data
parameters = load_job_data(rb_job)["parameters"]
exp_rb = InterleavedRB(
    interleaved_element=CXGate(),
    physical_qubits=[QUBIT_C, QUBIT_T],
    lengths=parameters["lengths"],
    num_samples=parameters["num_sampling"],
    backend=backend,
    seed=parameters["seed"])
exp_data = retrieve_expdata(rb_job, backend, exp=exp_rb)

exp_data.figure(0)