import os
from typing import Any, List, Dict, Union

import numpy as np
import matplotlib.pyplot as plt

from qiskit import IBMQ, QuantumCircuit, QuantumRegister, ClassicalRegister, quantum_info as qi
from qiskit.providers.ibmq import RunnerResult
from qiskit.result import marginal_counts
import qiskit.tools.jupyter

%matplotlib inline

import warnings
warnings.filterwarnings("ignore")


hub = 'ibm-q-community'
group = 'ieee-session'
project = 'event-2021'
backend_name = 'ibm_perth'

from qiskit import IBMQ
IBMQ.load_account()
provider = IBMQ.get_provider(hub=hub, group=group, project=project)
backend = provider.get_backend(backend_name)

import qiskit.tools.jupyter
backend

from qiskit import QuantumCircuit

# Fill in Readout test and draw

from qiskit import QuantumCircuit

qc_prep_excited = QuantumCircuit(1, 1)
qc_prep_excited.x(0)
qc_prep_excited.measure(0, 0)
qc_prep_excited.draw(output="mpl", idle_wires=False)

qubit = 3
shots = 1000

from qiskit import transpile

# Transpile for qubit and draw

from qiskit import transpile

qc_prep_excited_transpiled = transpile(qc_prep_excited, backend, initial_layout=[qubit])

qc_prep_excited_transpiled.draw(output="mpl", idle_wires=False)

from run_openqasm3 import run_openqasm3

# Run the circuit / QASM3

from run_openqasm3 import run_openqasm3

job = run_openqasm3(qc_prep_excited_transpiled, backend, verbose=True)

runtime_params = {
    'circuits': [qc_prep_excited_transpiled],
}

options = {'backend_name': backend.name()}

from qiskit.providers.ibmq import RunnerResult

prep_excited_job = provider.runtime.run(
    program_id="qasm3-runner",
    options=options,
    inputs=runtime_params,
    image="qiskit-program-qasm3:latest",
    result_decoder=RunnerResult,
)

prep_excited_result = prep_excited_job.result()
print(prep_excited_result.get_counts())

try:
    from qiskit import qasm3
except ImportError:
    print("You likely do not have the main branch of Qiskit installed. Please install it to run this code")
    raise
    


basis_gates = backend.configuration().basis_gates

def dump_qasm3(circuit, backend=backend):
    return qasm3.Exporter(includes=[], basis_gates=basis_gates).dumps(circuit)

excited_prep_source_string = dump_qasm3(qc_prep_excited_transpiled)
print(excited_prep_source_string)

runtime_params = {
    'circuits': [excited_prep_source_string],
}

options = {'backend_name': backend.name()}


prep_excited_job_source_str = provider.runtime.run(
    program_id="qasm3-runner",
    options=options,
    inputs=runtime_params,
    image="qiskit-program-qasm3:latest",
    result_decoder=RunnerResult,
)

print(prep_excited_job_source_str.result().get_counts())

qc_real_time = QuantumCircuit(1, 2)

# Program a builtin reset test circuit named job real time and draw


qc_real_time = QuantumCircuit(1, 2)
qc_real_time.x(0)
qc_real_time.measure(0, 0)
qc_real_time.reset(0)
qc_real_time.measure(0, 1)

qc_real_time.draw(output="mpl", idle_wires=False)

# Run the reset job and check the results
real_time_job = backend.run(transpile(qc_real_time, backend, initial_layout=[qubit]), shots=1000)
print(real_time_job.job_id())

print(f" Our reset counts were: {real_time_job.result().get_counts(0)}")

job_real_time = backend.run(qc_real_time)
print(f"Job id: {job_real_time.job_id()}")
result_real_time = job_real_time.result() 
print(f"Reset results: {result_real_time.get_counts(0)}")


from qiskit import QuantumRegister, ClassicalRegister

qr = QuantumRegister(1)
crx = ClassicalRegister(1, name="xresult")
crm = ClassicalRegister(1, name="measureresult")
qc_reset = QuantumCircuit(qr, crx, crm, name="Reset")

# Program a reset directly using control flow and draw


from qiskit import QuantumRegister, ClassicalRegister

qr = QuantumRegister(1)
crx = ClassicalRegister(1, name="xresult")
crm = ClassicalRegister(1, name="measureresult")
qc_reset = QuantumCircuit(qr, crx, crm, name="Reset")
qc_reset.x(0)
qc_reset.measure(0, crx)
qc_reset.x(0).c_if(crx, 1)
qc_reset.measure(0, crm)
 
qc_reset = transpile(qc_reset, backend, initial_layout=[1])
qc_reset.draw(output="mpl", idle_wires=False)

print(dump_qasm3(qc_reset))

reset_job = run_openqasm3(qc_reset, backend, verbose=True)


from qiskit.result import marginal_counts

reset_result = reset_job.result()
reset_counts = reset_result.get_counts(0)
mitigated_reset_results = marginal_counts(reset_counts, indices=[1])
print(f"Full counts including reset: {reset_counts}")
print(f"Results from our reset - |0>s prepared {mitigated_reset_results.get('0')}, |1>s prepared {mitigated_reset_results['1']}"
)
      


# Program several rounds of reset using QASM3

#print(dump_qasm3(qc_reset))

qubit_init_qasm3 = """

"""

qubit_init_qasm3 = """
    
OPENQASM 3;

bit reset_result0;
bit reset_result1;
bit reset_result2;
bit[1] measureresult;

x $1;
reset_result0 = measure $1;

// three rounds of reset.
if (reset_result0 == 1) {
    x $1;
}

reset_result1 = measure $1;

if (reset_result1 == 1) {
    x $1;
}

reset_result2 = measure $1;

if (reset_result2 == 1) {
    x $1;
}

// Final measurement
measureresult[0] = measure $1;

"""

init_job = run_openqasm3(qubit_init_qasm3, backend, verbose=False)
print(init_job.job_id())
better_reset_result_1s = sum(count for key, count in init_job.result().get_counts(0).items() if (len(key)==4))
print(f"Results from our better reset - |0>s prepared {shots-better_reset_result_1s}, |1>s prepared {better_reset_result_1s} ")


qr = QuantumRegister(3)
crz = ClassicalRegister(1, name="crz")
crx = ClassicalRegister(1, name="crx")
result = ClassicalRegister(1, name="result")

qc_teleport = QuantumCircuit(qr, crz, crx, result, name="Teleport")

# Apply teleportation circuit
qc_teleport.h(qr[1])
qc_teleport.cx(qr[1], qr[2])
#qc_teleport.barrier(qr)
qc_teleport.cx(qr[0], qr[1])
qc_teleport.h(qr[0])
qc_teleport.measure(qr[0], crz)
qc_teleport.measure(qr[1], crx)
qc_teleport.z(qr[2]).c_if(crz, 1)
qc_teleport.x(qr[2]).c_if(crx, 1)
qc_teleport.draw(output="mpl")

qc_state_prep = QuantumCircuit(1)
qc_state_prep.x(0)

print(qc_state_prep)

target_state_prep = qi.Statevector.from_instruction(qc_state_prep)
print(target_state_prep)

teleport_qubits = [0, 1, 2]

qc_teleport_state = QuantumCircuit(qr, crz, crx, result, name="Teleport Hadamard")

def init(qc, qr):
    qc.barrier(qr)
    qc.reset(qr)
    qc.barrier(qr)
    qc.reset(qr)
    qc.barrier(qr)
    qc.reset(qr)
    qc.barrier(qr)

init(qc_teleport_state, qr)

# Prepare state to teleport
qc_teleport_state.compose(qc_state_prep, [qr[0]], inplace=True)
qc_teleport_state.barrier(qr)

# Compose with teleportation circuit
qc_teleport_state.compose(qc_teleport, inplace=True)

qc_teleport_state.draw(output="mpl")

qc_teleport_experiment = qc_teleport_state.copy()
qc_teleport_experiment.measure(qr[2], result)
qc_teleport_experiment = transpile(qc_teleport_experiment, backend, initial_layout=teleport_qubits, optimization_level=3)
qc_teleport_experiment.draw(output="mpl")

teleport_job = run_openqasm3(qc_teleport_experiment, backend, verbose=False)
print(teleport_job.job_id())

teleport_result = teleport_job.result()
print(f"All teleportation reults: {teleport_result.get_counts(0)}")

marginal_teleport_counts = marginal_counts(teleport_result.get_counts(0), indices=[2])
print(f"Marginalized teleportation counts: {marginal_teleport_counts}")

from qiskit.ignis.verification.tomography import state_tomography_circuits, StateTomographyFitter
from qiskit.ignis.mitigation.measurement import complete_meas_cal, CompleteMeasFitter
from qiskit.visualization import plot_state_city, plot_state_paulivec

meas_calibs, state_labels = complete_meas_cal(qr=[qr[2]], circlabel='mcal')
meas_calibs = transpile(meas_calibs, backend, initial_layout={qr[2]: teleport_qubits[-1]})

meas_calibs[1].draw(idle_wires=False, output="mpl")

meas_calib_job = run_openqasm3(meas_calibs, backend, verbose=False)
print(f"Job id: {meas_calib_job.job_id()}")

meas_calib_result = meas_calib_job.result()

meas_fitter = CompleteMeasFitter(meas_calib_result, state_labels, circlabel='mcal')
meas_fitter.cal_matrix

meas_filter = meas_fitter.filter
mitigated_teleport_counts = meas_filter.apply(marginal_teleport_counts)

print(f"Unmitigated teleportation counts {marginal_teleport_counts}")
print(f"Mitigated teleportation counts {mitigated_teleport_counts}")

def build_state_tomo_circuits(qc_prep, qr, clbits, tomo_qubits, layout):
    qst_tomography_circuits = state_tomography_circuits(QuantumCircuit(qr), tomo_qubits)
    
    tomography_circuits = []
    for tomo_circuit in qst_tomography_circuits:
        tomo_circ = qc_prep.compose(tomo_circuit, clbits=clbits)
        tomo_circ.name = tomo_circuit.name
        tomography_circuits.append(tomo_circ)
        
    return transpile(tomography_circuits, backend, initial_layout=layout)

qt_qst_tomography_circuits = build_state_tomo_circuits(qc_teleport_state, qr, [result[0]], [qr[2]], teleport_qubits)


teleport_qst_jobs = run_openqasm3(qt_qst_tomography_circuits, backend, verbose=False)

qt_qst_tomography_circuits[1].draw(output="mpl")

teleport_qst_result = teleport_qst_jobs.result()
marginal_teleport_qst_result = marginal_counts(teleport_qst_result, indices=[2])
mitigated_qst_results = meas_filter.apply(marginal_teleport_qst_result)

qt_qst_tomo_fitter = StateTomographyFitter(mitigated_qst_results, qt_qst_tomography_circuits)

qt_qst_rho_fit = qt_qst_tomo_fitter.fit(method='lstsq')
qt_qst_F = qi.state_fidelity(qt_qst_rho_fit, target_state_prep)
print('State Fidelity: F = {:.5f}'.format(qt_qst_F))

fig, axes = plt.subplots(1, 2, sharey=True)
plot_state_paulivec(qt_qst_rho_fit, ax=axes[0]);
plot_state_paulivec(qi.DensityMatrix.from_instruction(qc_state_prep), ax=axes[1])
axes[1].set_ylabel(None);
axes[0].set_title("Experiment")
axes[1].set_title("Ideal")
plt.tight_layout()
fig.set_size_inches(14, 6)

plot_state_city(qt_qst_rho_fit, title="Experiment");

plot_state_city(qi.DensityMatrix.from_instruction(qc_state_prep), title="Ideal");