import numpy as np
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.test.mock import FakeRome
from qiskit.providers.aer import AerProvider
from qiskit.quantum_info import hellinger_fidelity


qr = QuantumRegister(2)
cr = ClassicalRegister(2)
qc = QuantumCircuit(qr, cr)

qc.h(0)
qc.cx(0,1)
qc.barrier()
qc.measure(qr, cr)

qc.draw('mpl')

num_shots = 4096

fake_rome = FakeRome()
job_fake_rome = fake_rome.run(qc, shots=4096)
result_fake_rome = job_fake_rome.result()
counts_fake_rome = result_fake_rome.get_counts()

print(counts_fake_rome)
pr_m = np.zeros(4, dtype=float)
for b, v in counts_fake_rome.items():
    pr_m[int(b, base=2)] = v/num_shots
print(pr_m)

pr_m_p = np.zeros((4,4), dtype=float)

qr_mit = QuantumRegister(2)
cr_mit = ClassicalRegister(2)

# state prep = 00
qc_00 = QuantumCircuit(qr_mit, cr_mit)
qc_00.barrier()
qc_00.measure(qr_mit, cr_mit)
counts_00 = fake_rome.run(qc_00, shots=num_shots)\
    .result().get_counts()

pr_m_p[0][0] = counts_00['00']/num_shots if '00' in counts_00 else 0.0
pr_m_p[1][0] = counts_00['01']/num_shots if '01' in counts_00 else 0.0
pr_m_p[2][0] = counts_00['10']/num_shots if '10' in counts_00 else 0.0
pr_m_p[3][0] = counts_00['11']/num_shots if '11' in counts_00 else 0.0


# state prep = 01
qc_01 = QuantumCircuit(qr_mit, cr_mit)

qc_01.x(0)

qc_01.barrier()
qc_01.measure(qr_mit, cr_mit)
counts_01 = fake_rome.run(qc_01, shots=num_shots)\
    .result().get_counts()

pr_m_p[0][1] = counts_01['00']/num_shots if '00' in counts_01 else 0.0
pr_m_p[1][1] = counts_01['01']/num_shots if '01' in counts_01 else 0.0
pr_m_p[2][1] = counts_01['10']/num_shots if '10' in counts_01 else 0.0
pr_m_p[3][1] = counts_01['11']/num_shots if '11' in counts_01 else 0.0


# state prep = 10
qc_10 = QuantumCircuit(qr_mit, cr_mit)

qc_10.x(1)

qc_10.barrier()
qc_10.measure(qr_mit, cr_mit)
counts_10 = fake_rome.run(qc_10, shots=num_shots)\
    .result().get_counts()

pr_m_p[0][2] = counts_10['00']/num_shots if '00' in counts_10 else 0.0
pr_m_p[1][2] = counts_10['01']/num_shots if '01' in counts_10 else 0.0
pr_m_p[2][2] = counts_10['10']/num_shots if '10' in counts_10 else 0.0
pr_m_p[3][2] = counts_10['11']/num_shots if '11' in counts_10 else 0.0


# state prep = 11
qc_11 = QuantumCircuit(qr_mit, cr_mit)

qc_11.x(0)
qc_11.x(1)

qc_11.barrier()
qc_11.measure(qr_mit, cr_mit)
counts_11 = fake_rome.run(qc_11, shots=num_shots)\
    .result().get_counts()

pr_m_p[0][3] = counts_11['00']/num_shots if '00' in counts_11 else 0.0
pr_m_p[1][3] = counts_11['01']/num_shots if '01' in counts_11 else 0.0
pr_m_p[2][3] = counts_11['10']/num_shots if '10' in counts_11 else 0.0
pr_m_p[3][3] = counts_11['11']/num_shots if '11' in counts_11 else 0.0

np.set_printoptions(suppress=True)
print(pr_m_p)

pr_p = np.matmul(np.linalg.inv(pr_m_p), pr_m)
print("Before mitigation: ", pr_m)
print("After  mitigation: ", pr_p)

qasm_backend = AerProvider().get_backend("qasm_simulator")
counts_qasm = qasm_backend.run(qc, shots=num_shots).result().get_counts()

pr_ideal = np.zeros(4, dtype=float)
for b, v in counts_qasm.items():
    pr_ideal[int(b, base=2)] = v/num_shots
print("Ideal probs: ", pr_ideal)

# convert pr_p to dict form
counts_mit = dict()
for i, v in enumerate(pr_p):
    counts_mit[bin(i)[2:].rjust(2, '0')] = v * num_shots

print("Fidelity without mitigation: ", hellinger_fidelity(counts_fake_rome, counts_qasm))
print("Fidelity with mitigation: ", hellinger_fidelity(counts_mit, counts_qasm))

from qiskit.visualization import plot_histogram
plot_histogram([counts_fake_rome, counts_mit, counts_qasm],
               legend=['no mit', 'mitigated', 'ideal'])

from qiskit.ignis.mitigation import CompleteMeasFitter, complete_meas_cal

meas_calibs, state_labels = complete_meas_cal(qr=qr_mit, circlabel='mcal')
for ckt in meas_calibs:
    print(ckt)

job_calibs = fake_rome.run(meas_calibs, shots=num_shots)
meas_fitter = CompleteMeasFitter(job_calibs.result(), state_labels, circlabel='mcal')

print(meas_fitter.cal_matrix)
meas_fitter.plot_calibration()

# Get the filter object
meas_filter = meas_fitter.filter

# Results with mitigation
mitigated_results = meas_filter.apply(result_fake_rome)
mitigated_counts = mitigated_results.get_counts(0)

print("Fidelity with mitigation: ", hellinger_fidelity(mitigated_counts, counts_qasm))