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

import numpy as np
import matplotlib.pyplot as plt

from qiskit import IBMQ, Aer, QuantumCircuit, QuantumRegister, ClassicalRegister, transpile, execute
from qiskit.tools.visualization import plot_histogram
from qiskit.result import marginal_counts
import qiskit.tools.jupyter

from run_openqasm3 import run_openqasm3

%matplotlib inline

import warnings
warnings.filterwarnings("ignore")

pi = np.pi

backend_sim = Aer.get_backend('aer_simulator')

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

IBMQ.load_account()
provider = IBMQ.get_provider(hub=hub, group=group, project=project)
backend_real = provider.get_backend(backend_name)

basis_gates = backend_real.configuration().basis_gates

import qiskit.tools.jupyter
backend_real

from qiskit import qasm3

def dump_qasm3(circuit, backend=backend_real):
    return qasm3.Exporter(includes=[], basis_gates=basis_gates+["reset"], disable_constants=True).dumps(circuit)

q = QuantumRegister(1,'q')
c = ClassicalRegister(1,'c')
qc = QuantumCircuit(q, c)
qc.h(0)
qc.measure(0,0)
qc.x(0).c_if(c, 0)
qc.draw(output='mpl')

q = QuantumRegister(3,'q')
c = ClassicalRegister(3,'c')
qc = QuantumCircuit(q, c)
qc.h(q[0])
qc.h(q[1])
qc.h(q[2])
qc.barrier()
qc.measure(q,c)
qc.draw('mpl')

print(bin(3))
print(bin(7))

q = QuantumRegister(3,'q')
c = ClassicalRegister(3,'c')
qc = QuantumCircuit(q, c)
qc.h(0)
qc.h(1)
qc.h(2)
qc.barrier()
qc.measure(q,c)

qc.x(2).c_if(c, 3) # for the 011 case
qc.x(2).c_if(c, 7) # for the 111 case

qc.draw(output='mpl')

qubits = [5, 6]

def reset_conditional(qc, q, reset_reg):
    qc.barrier(q)
    qc.measure(q, reset_reg)
    qc.x(q).c_if(reset_reg, 1)
    qc.barrier(q)

def reset_conditional(qc, q, reset_reg):
    qc.barrier(q)
    qc.measure(q, reset_reg)
    qc.x(q).c_if(reset_reg, 1)
    qc.barrier(q)

qr = QuantumRegister(2)
reset_reg0 = ClassicalRegister(1, "reset_reg0")
reset_reg1 = ClassicalRegister(1, "reset_reg1")
cr0 = ClassicalRegister(1, "reset_result0")
cr1 = ClassicalRegister(1, "reset_result1")

qc_reset = QuantumCircuit(qr, reset_reg0, reset_reg1, cr0, cr1)
reset_conditional(qc_reset, qr[0], reset_reg0)
reset_conditional(qc_reset, qr[1], reset_reg1)

qc_reset.measure(0, cr0)
qc_reset.measure(1, cr1)

qc_reset = transpile(qc_reset, backend_real, initial_layout=qubits)
qc_reset.draw(output="mpl")

# Hint use qiskit.result.marginal_counts to remove the first three classical register results
# from your counts
def pad_counts(counts):
    new_counts = {}
    for k, v in counts.items():
        if len(k) < 4:
            for _ in range(4-len(k)):
                k = "0"+k
        new_counts[k] = v
    return new_counts

def remove_reset_results(counts):
    return marginal_counts(pad_counts(counts), indices=[0, 1])



reset_job = run_openqasm3(qc_reset, backend_real, verbose=False)
print(f"Reset job id: {reset_job.job_id()}")
reset_result = reset_job.result()
reset_counts = remove_reset_results(reset_result.get_counts(0))
print(f"The results of your reset circuit are {reset_counts}")
assert reset_counts["00"] > 700

def init(qc, qr):
    """Implement a qubit initialization using three rounds of qubit reset."""
    qc.barrier(qr)
    qc.reset(qr)
    qc.reset(qr)
    qc.reset(qr)
    qc.barrier(qr)
    

    
qr = QuantumRegister(2)
cr0 = ClassicalRegister(1, "reset_result0")
cr1 = ClassicalRegister(1, "reset_result1")

qc_init = QuantumCircuit(qr, cr0, cr1)

init(qc_init, qr)

qc_init.measure(0, cr0)
qc_init.measure(1, cr1)

qc_init = transpile(qc_init, backend_real, initial_layout=qubits)
qc_init.draw(output="mpl")

init_job = run_openqasm3(qc_init, backend_real, verbose=False)
print(f"Qubit init job id: {init_job.job_id()}")
init_result = init_job.result()
init_counts = init_result.get_counts()
print(f"The results of your reset circuit are {init_counts}")
assert init_counts["0 0"] > 900

delta_counts = init_counts["0 0"] - reset_counts["00"]
print(f"With the initialization circuit we prepared the ground state {delta_counts} more times than in the reset circuit.""")

nq = 2
m = 2

qr = QuantumRegister(2,'q')
c0 = ClassicalRegister(1,'c0')
c1 = ClassicalRegister(1, 'c1')
qc_S = QuantumCircuit(qr, c0, c1)

def init_ipe_circuit(qc):
    """Initialize the qubits in your IPE circuit."""
    init(qc, [0, 1])

init_ipe_circuit(qc_S)
qc_S.draw(output="mpl")

def initialize_eigenstate(qc):
    """Initialize the eigenstate and prepare our ancilla qubit"""
    qc.h(0)
    qc.x(1)

initialize_eigenstate(qc_S)
    
qc_S.draw('mpl')


theta = 1 * np.pi / 2


cu_circ = QuantumCircuit(2)
cu_circ.cp(theta,0,1)
cu_circ.draw('mpl')


for _ in range(2**(m-1)):
    qc_S.cp(theta,0,1)
qc_S.draw('mpl')

def x_measurement(qc, qubit, cbit):
    """Measure 'qubit' in the X-basis, and store the result in 'cbit'"""
    qc.h(qubit)
    qc.measure(qubit, cbit)
    
x_measurement(qc_S, qr[0], c0)
qc_S.draw('mpl')


def reset_auxiliary(qc):
    qc.reset(0)
    qc.h(0)

reset_auxiliary(qc_S)
qc_S.draw('mpl')

qc_S.p(-np.pi/2,0).c_if(c0, 1)
qc_S.draw('mpl')

## 2^t c-U operations (with t=m-2)
for _ in range(2**(m-2)):
    qc_S.cp(theta,0,1)

x_measurement(qc_S, qr[0], c1) 

qc_S.draw('mpl')

count = execute(qc_S, backend_sim).result().get_counts()
count

count0 = remove_reset_results(execute(qc_S, backend_sim).result().get_counts())

key_new = [str(int(key,2)/2**m) for key in list(count0.keys())]
count1 = dict(zip(key_new, count0.values()))

fig, ax = plt.subplots(1,2)
plot_histogram(count0, ax=ax[0])
plot_histogram(count1, ax=ax[1])
plt.tight_layout()

qc_S_backend = transpile(qc_S, backend_real, initial_layout=qubits, optimization_level=3)
qc_S_backend.draw(output="mpl")

real_ipe_job = run_openqasm3(qc_S_backend, backend_real, verbose=False)
print(f"IPE job id: {real_ipe_job.job_id()}")

real_counts0 = remove_reset_results(real_ipe_job.result().get_counts())
real_key_new = [str(int(key,2)/2**m) for key in list(real_counts0.keys())]
real_counts1 = dict(zip(real_key_new, real_counts0.values()))

fig, ax = plt.subplots(1,2)
plot_histogram(real_counts0, ax=ax[0])
plot_histogram(real_counts1, ax=ax[1])
plt.tight_layout()

import qiskit.tools.jupyter
%qiskit_version_table
%qiskit_copyright