from qiskit.circuit.random import random_circuit

circuit = random_circuit(2, 4, seed=1).decompose(reps=1)
display(circuit.draw("mpl"))

from qiskit.quantum_info import SparsePauliOp

observable = SparsePauliOp("ZZ")

from qiskit.primitives import Estimator

estimator = Estimator()
job = estimator.run(circuit, observable)
simulation = job.result()

from zne.utils.serialization import EstimatorResultEncoder

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

from qiskit.primitives import BackendEstimator
from qiskit.providers.fake_provider import FakeNairobi  # We need noise!
from zne import zne

ZNEEstimator = zne(BackendEstimator)  # Any implementation of BaseEstimator is valid
backend = FakeNairobi()
estimator = ZNEEstimator(backend=backend)

from zne import ZNEStrategy
from zne.noise_amplification import LocalFoldingAmplifier
from zne.extrapolation import PolynomialExtrapolator

## NOISE FACTORS
noise_factors = range(1,9)

## NOISE AMPLIFIER
noise_amplifier = LocalFoldingAmplifier(
    gates_to_fold=2,
)

## EXTRPOLATOR
extrapolator = PolynomialExtrapolator(
    degree=2,
)

## ZNE STRATEGY
zne_strategy = ZNEStrategy(
    noise_factors=noise_factors,
    noise_amplifier=noise_amplifier,
    extrapolator=extrapolator,
)

job = estimator.run(circuit, observable, zne_strategy=zne_strategy)
result = job.result()

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

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"})

############################  DATA  ############################
exact = simulation.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()

circuit = random_circuit(2, 2, seed=0).decompose(reps=1)
display(circuit.draw("mpl"))

noise_amplifier = LocalFoldingAmplifier(
    gates_to_fold = 2,
    sub_folding_option = "from_first",
    random_seed = None,
)

noisy_circuit = noise_amplifier.amplify_circuit_noise(circuit=circuit, noise_factor=2)
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(),
)
    
job = estimator.run(circuit, observable, zne_strategy=zne_strategy)    
result = job.result()

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

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

%qiskit_version_table
%qiskit_copyright