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 qiskit.circuit.random import random_circuit
from qiskit.quantum_info import SparsePauliOp

circuit = random_circuit(2, 4, seed=1).decompose(reps=1)
observable = SparsePauliOp("ZZ")

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

display(circuit.draw("mpl"))
print(f"  > Observable: {observable.paulis}")
print(f"  > Expectation value: {result.values[0]}")
print(f"  > Metadata: {result.metadata[0]}")

from qiskit.primitives import Estimator

circuit = random_circuit(2, 4, seed=1).decompose(reps=1)
observable = SparsePauliOp("ZZ")

job = Estimator().run(circuit, observable)  # !!!
result = job.result()

display(circuit.draw("mpl"))
print(f"  > Observable: {observable.paulis}")
print(f"  > Expectation value: {result.values[0]}")
print(f"  > Metadata: {result.metadata[0]}")

from zne import ZNEStrategy

circuit = random_circuit(2, 4, seed=1).decompose(reps=1)
observable = SparsePauliOp("ZZ")

zne_strategy = ZNEStrategy(noise_factors=[1, 3, 5])
job = estimator.run(circuit, observable, zne_strategy=zne_strategy)  # !!!
result = job.result()

display(circuit.draw("mpl"))
print(f"  > Observable: {observable.paulis}")
print(f"  > Expectation value: {result.values[0]}")
print(f"  > Metadata: {result.metadata[0]}")

from zne.utils.serialization import EstimatorResultEncoder

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

from zne.noise_amplification import LocalFoldingAmplifier, GlobalFoldingAmplifier

global_amplifier = GlobalFoldingAmplifier()
local_amplifier = LocalFoldingAmplifier(gates_to_fold=2)  # Fold two-qubit gates

zne_strategy = ZNEStrategy(
    noise_factors=range(1, 9), 
    noise_amplifier=local_amplifier,
)

from zne.extrapolation import PolynomialExtrapolator

extrapolator = PolynomialExtrapolator(degree=2)

zne_strategy = ZNEStrategy(
    noise_factors=[1.0, 1.5, 2.0], 
    extrapolator=extrapolator
)

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

circuit = random_circuit(2, 2, seed=0).decompose(reps=1)
observable = SparsePauliOp("ZZ")

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

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

display(circuit.draw("mpl"))
print(f"  > Observable: {observable.paulis}")
print(f"  > Expectation value: {result.values[0]}")
print(f"  > Metadata: {result.metadata[0]}")

from qiskit.circuit.random import random_circuit
from qiskit.primitives import BackendEstimator
from qiskit.providers.fake_provider import FakeNairobi
from qiskit.quantum_info import SparsePauliOp
from zne import zne, ZNEStrategy
from zne.noise_amplification import LocalFoldingAmplifier
from zne.extrapolation import PolynomialExtrapolator

## Build ZNEEstimator
ZNEEstimator = zne(BackendEstimator)
backend = FakeNairobi()
estimator = ZNEEstimator(backend=backend)

## Build our input circuit and observable
circuit = random_circuit(2, 4, seed=1).decompose(reps=1)
observable = SparsePauliOp("ZZ")

## Build a ZNE strategy
zne_strategy = ZNEStrategy(
    noise_factors=range(1, 9),
    noise_amplifier=LocalFoldingAmplifier(gates_to_fold=2),
    extrapolator=PolynomialExtrapolator(degree=2),
)

## Run experiment
job = estimator.run(circuit, observable, zne_strategy=zne_strategy)  # !!!
result = job.result()

## Display results
display(circuit.draw("mpl"))
print(f"  > Observable: {observable.paulis}")
print(f"  > Expectation value: {result.values[0]}")
print(f"  > Metadata: {result.metadata[0]}")

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 = 1  # From simulation above
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