import scipy as sp
import numpy as np
import matplotlib.pyplot as plt
import qiskit

from qiskit import QuantumCircuit, transpile, Aer, IBMQ
from qiskit.providers.aer.noise import NoiseModel
from qiskit.utils import QuantumInstance

# Loading your IBM Q account(s)
provider = IBMQ.load_account()

def run_obs(qc, shots, backend_name='ibmq_santiago'):
    """Helper function to get the expected value of the
       | 11 > state.
    """
    
    # Get device noise model
    device = provider.get_backend(backend_name)
    noise_model = NoiseModel.from_backend(device)
    coupling_map = device.configuration().coupling_map
    seed = 42

    # Define the backend
    backend = QuantumInstance(
                            backend=Aer.get_backend("qasm_simulator"),
                            seed_transpiler=seed,
                            optimization_level=1,
                            noise_model=noise_model,
                            shots=shots,
                            )
    qc = qc.copy()
    qc.measure_all()
    counts = backend.execute(qc).get_counts()
    return counts['11']/shots

qc = QuantumCircuit(2)
qc.x(0)
qc.cx(0, 1)
qc.draw('mpl')

run_obs(qc, 1000)

def fold_cx(qc, alpha):
    """ Fold the cx on the circuit given an alpha value.
    """
    d = qc.depth()
    k = np.ceil(d*(alpha - 1)/2)
    n = k//d
    s = k%d
    
    instructions = []
    for instruction, qargs, cargs in qc:
        if instruction.name == 'cx':
            instruction = qiskit.circuit.library.CXGate()
            barrier = qiskit.circuit.library.Barrier(len(qc))
            
            instructions.append((instruction, qargs, cargs))
            instructions.append((barrier, qargs, cargs))
            for _ in range(2*int(n + s)):
                instructions.append((instruction, qargs, cargs))
                instructions.append((barrier, qargs, cargs))
                
        
        else:
            instructions.append((instruction, qargs, cargs))
    
    folded_qc = qc.copy()
    folded_qc.data = instructions
    return folded_qc

alpha_list = [1, 2, 5, 7, 9, 11]
folds = []
for alpha in alpha_list:
    folds.append(fold_cx(qc, alpha))

expected_vals = []
for i in range(len(alpha_list)):
    expected_vals.append(run_obs(folds[i], shots=1000))

print(f"Alphas: {alpha_list}\nExpected values{expected_vals}")

plt.plot(alpha_list, expected_vals, 'o')
plt.xlabel(r"$\alpha$", size=14)
plt.ylabel(r"E($\alpha$)", size=14)
plt.show()

def linear(x, a, b):
    """Linear fit
    """
    return a*x + b

y = expected_vals
x = alpha_list
x, y = np.array(x), np.array(y)

popt, _ = sp.optimize.curve_fit(linear, x, y)
print("Mitigated expectation value:", np.round(popt[1],3))
print("Unmitigated expectation value:", expected_vals[0])

print("Absolute error with mitigation:", np.round(np.abs(1 - popt[1]), 3))
print("Absolute error without mitigation:", np.round(np.abs(1 - expected_vals[0]), 3))

plt.plot(alpha_list, expected_vals, 'o')
plt.plot(0, popt[1], '*', label="mitigated")
plt.xlabel(r"$\alpha$", size=14)
plt.ylabel(r"E($\alpha$)", size=14)
plt.legend()
plt.show()

import qiskit.tools.jupyter
%qiskit_version_table