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 = [1, 2]

def reset_conditional(qc, q, reset_reg):
    """Implement the conditional reset on the input circuit, qubit and classical register using c_if."""
    qc.barrier(q)
    # Add reset here
    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(0, 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, q):
    """Implement a qubit initialization using three rounds of qubit reset."""
    qc.barrier(q)
    # Add init here
    qc.barrier(q)
    

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

qc_init = QuantumCircuit(qr, cr0, cr1)
init(qc_init, qr[0])
init(qc_init, qr[1])

qc_init.measure(0, cr0)
qc_init.measure(0, 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

# Compare your counts here.

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

def init_ipe_circuit(qc):
    """Initialize the qubits in your IPE circuit."""
    qc.barrier(0, 1)
    # Place your code here
    qc.barrier(0, 1)

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

def initialize_eigenstate(qc):
    """Initialize the eigenstate and prepare our ancilla qubit"""
    pass

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


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


theta = 1 * np.pi / 2

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'"""
    pass
    
x_measurement(qc_S, q[0], c0)
qc_S.draw('mpl')


def reset_auxiliary(qc):
    """Reset your auxiliary qubit and then reinitialize it to |+>."""
    pass

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

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

nq = 2
m = 2

## 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, q[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)


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