## QISKIT
from qiskit.circuit import QuantumCircuit, Parameter
from qiskit.quantum_info import SparsePauliOp
from qiskit.primitives import Estimator, BackendEstimator
from qiskit.providers.fake_provider import FakeNairobi as Backend

## ZNE
from zne import zne, ZNEStrategy
from zne.noise_amplification import LocalFoldingAmplifier
from zne.extrapolation import PolynomialExtrapolator

## OTHER
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

mpl.rcParams["figure.dpi"] = 300
plt.rcParams.update({"text.usetex": True, "font.family": "Computer Modern Roman"})
plt.rcParams.update({'font.size': 16})

## CIRCUIT
cir = QuantumCircuit(2)
cir.h(0)
cir.cx(0, 1)
cir.ry(Parameter("theta"), 0)
display(cir.draw("mpl"))

## THETA
theta_vec = np.linspace(-np.pi, np.pi, 15)

## OBSERVABLE
obs1 = SparsePauliOp(["ZZ", "ZX", "XZ", "XX"], [1, -1, +1, 1])
obs2 = SparsePauliOp(["ZZ", "ZX", "XZ", "XX"], [1, +1, -1, 1])

## IDEAL ESTIMATOR
estimator = Estimator()

## JOBS
job1 = estimator.run(
    circuits = [cir] * len(theta_vec), 
    observables = [obs1] * len(theta_vec), 
    parameter_values = [[v] for v in theta_vec],
)
job2 = estimator.run(
    circuits = [cir] * len(theta_vec), 
    observables = [obs2] * len(theta_vec), 
    parameter_values = [[v] for v in theta_vec],
)


## RESULTS
CHSH1_ideal = job1.result()
CHSH2_ideal = job2.result()

## NOISY ESTIMATOR
backend = Backend()
estimator = BackendEstimator(backend=backend)

## JOBS
job1 = estimator.run(
    circuits = [cir] * len(theta_vec), 
    observables = [obs1] * len(theta_vec), 
    parameter_values = [[v] for v in theta_vec],
)
job2 = estimator.run(
    circuits = [cir] * len(theta_vec), 
    observables = [obs2] * len(theta_vec), 
    parameter_values = [[v] for v in theta_vec],
)


## RESULTS
CHSH1_noisy = job1.result()
CHSH2_noisy = job2.result()

## MITIGATED ESTIMATOR
ZNEEstimator = zne(BackendEstimator)
zne_strategy = ZNEStrategy(
    noise_factors=(1,7,13),
    noise_amplifier=LocalFoldingAmplifier()
)
estimator = ZNEEstimator(backend=backend, zne_strategy=zne_strategy)

## JOBS
job1 = estimator.run(
    circuits = [cir] * len(theta_vec), 
    observables = [obs1] * len(theta_vec), 
    parameter_values = [[v] for v in theta_vec],
)
job2 = estimator.run(
    circuits = [cir] * len(theta_vec), 
    observables = [obs2] * len(theta_vec), 
    parameter_values = [[v] for v in theta_vec],
)


## RESULTS
CHSH1_mitigated = job1.result()
CHSH2_mitigated = job2.result()

## CONFIG
# plt.rcParams.update({'font.size': 22, 'lines.markersize': 10, 'lines.linewidth': 1.8})
plt.figure(figsize=(8,5))

## LIMITS
plt.axhline(y=2, color='r', linestyle='-')
plt.axhline(y=-2, color='r', linestyle='-')
plt.axhline(y=np.sqrt(2)*2, color='k', linestyle='-.')
plt.axhline(y=-np.sqrt(2)*2, color='k', linestyle='-.')

## PLOTS
plt.plot(theta_vec,CHSH1_ideal,'s-',label = 'CHSH1 Ideal')
plt.plot(theta_vec,CHSH2_ideal,'s-',label = 'CHSH2 Ideal')
plt.plot(theta_vec,CHSH1_mitigated,'^-',label = 'CHSH1 Mitigated')
plt.plot(theta_vec,CHSH2_mitigated,'^-',label = 'CHSH2 Mitigated')
plt.plot(theta_vec,CHSH1_noisy,'o-',label = 'CHSH1 Noisy')
plt.plot(theta_vec,CHSH2_noisy,'o-',label = 'CHSH2 Noisy')

## EMBELLISH 
plt.title("Zero noise extrapolation on CHSH")
plt.xlabel('Theta')
plt.ylabel('CHSH witness')
plt.grid(which='major',axis='both')
plt.rcParams.update({'font.size': 11})
plt.legend()

plt.show()

print(" # |  MITIGATED   <-    (NOISY DATA POINTS) \n" + "-" * 70)
for i in range(len(CHSH1_mitigated.values)):
    print(
        f"{i:2} |  "
        f"{CHSH1_mitigated.values[i]:+.5f}"
        f"    <-    "
        f"{CHSH1_mitigated.metadata[i]['zne']['noise_amplification']['values']}"
    )

import qiskit.tools.jupyter

%qiskit_version_table
%qiskit_copyright