from warnings import filterwarnings
filterwarnings('ignore')

import json
import numpy.random as npr
import matplotlib as mpl
import matplotlib.pyplot as plt
%matplotlib inline
# mpl.rcParams["figure.dpi"] = 300
# plt.rcParams.update({"text.usetex": True, "font.family": "Computer Modern Roman"})

from qiskit.circuit.library import TwoLocal
from qiskit.circuit.random import random_circuit

circuits = [
    TwoLocal(2, rotation_blocks=["ry", "rz"], entanglement_blocks="cz").decompose(reps=1),
    random_circuit(2, 4, seed=1).decompose(reps=1),
    random_circuit(2, 2, seed=0).decompose(reps=1)
]

_ = [display(c.draw("mpl")) for c in circuits]

from qiskit.quantum_info import SparsePauliOp

observables = [
    SparsePauliOp("ZZ"),
    SparsePauliOp("XX"),
    SparsePauliOp("YY"),
]

from qiskit.primitives import Estimator, BackendEstimator
from qiskit.providers.fake_provider import FakeNairobi as Backend


estimator = BackendEstimator(backend=Backend())

# measure expectation values of observable-0 and 1 on circuit-0 (with specified parameters)
job_0 = estimator.run(
    circuits=[circuits[0]] * len(observables),
    observables=observables,
    parameter_values=[[*npr.random(16)]] * len(observables),
    shots=2048,
)

# measure expectation values of observable-0 and 1 on circuit-1
job_1 = estimator.run(
    circuits=[circuits[1]] * len(observables),
    observables=observables,
    shots=2048,
)

result_1 = job_1.result()
result_1

result_1  # print again for comparison on consecutive runs

from zne import zne
ZNEEstimator = zne(BackendEstimator)  # Any implementation of BaseEstimator can be passed

from zne import ZNEStrategy
# from zne.noise_amplification import GlobalFoldingAmplifier, LocalFoldingAmplifier
# from zne.extrapolation import LinearExtrapolator, PolynomialExtrapolator

zne_strategy = ZNEStrategy(
    noise_factors=(1, 3, 5),
    # noise_amplifier = LocalFoldingAmplifier()
    # extrapolator = LinearExtrapolator()
)

estimator = ZNEEstimator(
    backend=Backend(), 
    zne_strategy=zne_strategy  # !!!
)

# measure expectation values of observable-0 and 1 on circuit-0 (with specified parameters)
job_0 = estimator.run(
    circuits=[circuits[0]] * len(observables),
    observables=observables,
    parameter_values=[[*npr.random(16)]] * len(observables),
    shots=2048,
)

# measure expectation values of observable-0 and 1 on circuit-1
job_1 = estimator.run(
    circuits=[circuits[1]] * len(observables),
    observables=observables,
    shots=2048,
)

result_1 = job_1.result()
result_1

from zne.utils.serialization import EstimatorResultEncoder

print(EstimatorResultEncoder.dumps(result_1, indent=2))

from zne import NOISE_AMPLIFIER_LIBRARY
NOISE_AMPLIFIER_LIBRARY.keys()

from zne.noise_amplification import GlobalFoldingAmplifier, LocalFoldingAmplifier

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

noise_amplifier = GlobalFoldingAmplifier()

noisy_circuit = noise_amplifier.amplify_circuit_noise(circuit=circuits[2], noise_factor=1)

display(noisy_circuit.draw("mpl", plot_barriers=True))

noise_amplifier = LocalFoldingAmplifier()

noisy_circuit = noise_amplifier.amplify_circuit_noise(circuit=circuits[2], noise_factor=1)

display(noisy_circuit.draw("mpl", plot_barriers=False))

noise_amplifier = GlobalFoldingAmplifier(
    sub_folding_option = "from_first",
    # random_seed = 123,
)

noisy_circuit = noise_amplifier.amplify_circuit_noise(circuit=circuits[2], noise_factor=4)

display(noisy_circuit.draw("mpl", plot_barriers=True))

noise_amplifier = LocalFoldingAmplifier(
    sub_folding_option = "from_first",
    # random_seed = 123,
)

noisy_circuit = noise_amplifier.amplify_circuit_noise(circuit=circuits[2], noise_factor=4)

display(noisy_circuit.draw("mpl", plot_barriers=False))

noise_amplifier = LocalFoldingAmplifier(
    gates_to_fold = 2,  # Two-qubit gates
    # gates_to_fold = "cx",
    # gates_to_fold = ["cx", "rz"],
    # gates_to_fold = [2, "rz"],
)

noisy_circuit = noise_amplifier.amplify_circuit_noise(circuit=circuits[2], noise_factor=3)

display(noisy_circuit.draw("mpl", plot_barriers=False))

from zne.extrapolation import Extrapolator, ReckoningResult
from zne.noise_amplification import NoiseAmplifier


############################  NOISE AMPLIFIER  ############################
class CustomAmplifier(NoiseAmplifier):
    def amplify_circuit_noise(self, circuit, noise_factor):
        return circuit.copy()  # Dummy, nonperforming
    
    def amplify_dag_noise(self, dag, noise_factor):
        return super().amplify_dag_noise(dag, noise_factor)


############################  EXTRAPOLATOR  ############################
class CustomExtrapolator(Extrapolator):
    @property
    def min_points(self):
        return 2
    
    def _extrapolate_zero(self, x_data, y_data, sigma_x, sigma_y):
        value = 1.0
        std_error = 1.0
        metadata = {"meta": "data"}
        return ReckoningResult(value, std_error, metadata)  # Dummy, nonperforming

zne_strategy = ZNEStrategy(
    noise_amplifier=CustomAmplifier(),
    noise_factors=(1, 3),
    extrapolator=CustomExtrapolator(),
)

estimator = ZNEEstimator(backend=Backend(), zne_strategy=zne_strategy)
    
job = estimator.run(
    circuits=circuits[1],
    observables=observables[1],
)
    
result = job.result()
print(EstimatorResultEncoder.dumps(result, indent=2))

cir = circuits[1]
display(cir.draw("mpl"))

obs = observables[0]
print(obs)

noise_factors = range(1,9)
list(noise_factors)

from zne.extrapolation import PolynomialExtrapolator

## ZNE STRATEGY
noise_amplifier = LocalFoldingAmplifier(
    gates_to_fold=2,
    sub_folding_option='from_first',
)

extrapolator = PolynomialExtrapolator(
    degree=2,
)

zne_strategy = ZNEStrategy(
    noise_factors=noise_factors,
    noise_amplifier=noise_amplifier,
    extrapolator=extrapolator,
)

## CALCULATION
estimator = ZNEEstimator(backend=Backend(), zne_strategy=zne_strategy)
    
job = estimator.run(
    circuits=cir,
    observables=obs,
)
    
result = job.result()

############################  DATA  ############################
exact = Estimator().run(cir, obs).result().values[0]
mitigated = result.values[0]
noise_factors = result.metadata[0]["zne"]["noise_amplification"]["noise_factors"]
noisy_values = result.metadata[0]["zne"]["noise_amplification"]["values"]

############################  PLOT  ############################
plt.rcParams["figure.figsize"] = (5,3.5)
plt.grid(which='major',axis='both')

plt.plot([0, noise_factors[-1]], [exact, exact], "--", label=f"Exact", color="#000000")
plt.scatter(0, mitigated, label=f"Mitigated", marker="x", color="#785ef0")
if noise_factors[0] == 1:
    plt.scatter(
        noise_factors[0], noisy_values[0], 
        label=f"Unmitigated", marker="s", color="#dc267f",
    )
plt.plot(
    noise_factors, noisy_values, 
    label=f"Amplified", marker="o", color="#dc267f",
)

plt.title("Zero noise extrapolation")
plt.xlabel("Noise Factor ($n$)")
plt.ylabel(f"Expectation Value ($\langle ZZ \\rangle$)")
plt.legend()

plt.show()

import qiskit.tools.jupyter  # pylint: disable=unused-import,wrong-import-order

%qiskit_version_table
%qiskit_copyright