# for loading execution context parameters from the environment
import os
from dotenv import load_dotenv

# for (de)serializing results
import json

# for processing data
import numpy as np

# for plotting results
import matplotlib as mpl
import matplotlib.pyplot as plt

# for defining the circuits and observables constituting our expectation value problems
from qiskit import QuantumCircuit
from qiskit.circuit import Parameter
from qiskit.quantum_info import SparsePauliOp

# for using real quantum hardware
from qiskit import IBMQ, transpile
from qiskit_ibm_runtime import QiskitRuntimeService

from zne import ZNEStrategy
from zne import Estimator as ZneEstimator
from zne.noise_amplification import NoiseAmplifier

load_dotenv()

TOKEN = os.environ["TOKEN"]

HUB = os.environ["HUB"]
GROUP = os.environ["GROUP"]
PROJECT = os.environ["PROJECT"]

BACKEND = os.environ["BACKEND"]

service = QiskitRuntimeService(channel="ibm_quantum", instance=f"{HUB}/{GROUP}/{PROJECT}", token=TOKEN)

IBMQ.enable_account(TOKEN)
provider = IBMQ.get_provider(hub=HUB, group=GROUP, project=PROJECT)
backend = provider.get_backend(BACKEND)

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

def plot_trial(
    noise_factors,
    measured_evs,
    measured_errors,
    trial_id="A",
    alpha=0.5,
    color="#dc267f",
):
    plt.plot(noise_factors, measured_evs, label=f"Measured (Trial {trial_id})", color=color)
    plt.fill_between(
        noise_factors,
        measured_evs - measured_errors,
        measured_evs + measured_errors,
        alpha=alpha,
        color=color,
        linewidth=0
    )

def show_experiment(
    trials,
    exact_ev,
    experiment_repr,
    observable_repr,
    num_shots=4000,
    ev_range=None
):
    noise_factors = np.array(trials[0]["noise_factors"])
    plt.plot([0, noise_factors[-1]], [exact_ev, exact_ev], "--", label=f"Exact", color="#000000")
    for trial_id, color, measured_data in zip(('A', 'B'), ("#dc267f", "#785ef0"), trials):
        measured_evs = np.array(measured_data["values"])
        measured_errors = np.sqrt(np.array(measured_data["variances"]) / num_shots)
        plot_trial(
            noise_factors,
            measured_evs,
            measured_errors,
            trial_id=trial_id,
            color=color
        )
    plt.title(experiment_repr)
    plt.xlabel("Noise Factor ($n$)")
    plt.ylabel(f"Expectation Value ($\langle {observable_repr} \\rangle$)")
    if ev_range: 
        plt.ylim(ev_range)
    plt.legend()
    plt.show()

class SingleGateFoldingAmplifier(NoiseAmplifier):
    def __init__(self, gate_name: str):
        self.gate_name = gate_name

    def amplify_circuit_noise(self, circuit: QuantumCircuit, noise_factor: float) -> QuantumCircuit:
        noisy_circuit = circuit.copy_empty_like()
        for instruction, qargs, cargs in circuit:
            if instruction.name == self.gate_name:
                for i in range(noise_factor):
                    if i % 2 == 0:
                        noisy_circuit.append(instruction, qargs, cargs)
                    else:
                        noisy_circuit.append(instruction.inverse(), qargs, cargs)
                    noisy_circuit.barrier(qargs)
            else:
                noisy_circuit.append(instruction, qargs, cargs)
        return noisy_circuit

x_circuit = QuantumCircuit(1)
x_circuit.x(0)

x_observable = SparsePauliOp.from_list([('Z', 1)])

x_exact_ev = -1.0

display(x_circuit.draw("mpl"))
print(x_observable)
print(x_exact_ev)

x_granular_noise_factors = tuple(range(1, 599 + 2, 2))
x_behavioral_noise_factors = tuple(range(1, 5999, 20))

# with ZneEstimator(
#     circuits=[x_circuit],
#     observables=[x_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=x_granular_noise_factors, noise_amplifier=SingleGateFoldingAmplifier()),
# ) as estimator:
#     x_granular_a_results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# x_granular_a_data = {
#     "noise_factors": list(x_granular_a_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(x_granular_a_results.metadata[0]["zne"]["values"]),
#     "variances": list(x_granular_a_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/x/x_granular_a.json', 'w') as f:
#     json.dump(x_granular_a_data, f)

with open('data/x/x_granular_a.json', 'r') as f:
    x_granular_a_data = json.load(f)

# with ZneEstimator(
#     circuits=[x_circuit],
#     observables=[x_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=x_granular_noise_factors, noise_amplifier=SingleGateFoldingAmplifier()),
# ) as estimator:
#     x_granular_b_results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# x_granular_b_data = {
#     "noise_factors": list(x_granular_b_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(x_granular_b_results.metadata[0]["zne"]["values"]),
#     "variances": list(x_granular_b_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/x/x_granular_b.json', 'w') as f:
#     json.dump(x_granular_b_data, f)

with open('data/x/x_granular_b.json', 'r') as f:
    x_granular_b_data = json.load(f)

# with ZneEstimator(
#     circuits=[x_circuit],
#     observables=[x_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=x_behavioral_noise_factors, noise_amplifier=SingleGateFoldingAmplifier()),
# ) as estimator:
#     x_behavioral_a_results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# x_behavioral_a_data = {
#     "noise_factors": list(x_behavioral_a_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(x_behavioral_a_results.metadata[0]["zne"]["values"]),
#     "variances": list(x_behavioral_a_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/x/x_behavioral_a.json', 'w') as f:
#     json.dump(x_behavioral_a_data, f)

with open('data/x/x_behavioral_a.json', 'r') as f:
    x_behavioral_a_data = json.load(f)

# with ZneEstimator(
#     circuits=[x_circuit],
#     observables=[x_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=x_behavioral_noise_factors, noise_amplifier=SingleGateFoldingAmplifier()),
# ) as estimator:
#     x_behavioral_b_results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# x_behavioral_b_data = {
#     "noise_factors": list(x_behavioral_b_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(x_behavioral_b_results.metadata[0]["zne"]["values"]),
#     "variances": list(x_behavioral_b_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/x/x_behavioral_b.json', 'w') as f:
#     json.dump(x_behavioral_b_data, f)

with open('data/x/x_behavioral_b.json', 'r') as f:
    x_behavioral_b_data = json.load(f)

show_experiment([x_granular_a_data, x_granular_b_data], x_exact_ev, "$X$ Granular Folding Profile — $\langle 0 | X^n Z X^n | 0 \\rangle(n)$", 'Z')

show_experiment([x_behavioral_a_data, x_behavioral_b_data], x_exact_ev, "$X$ Behavioral Folding Profile — $\langle 0 | X^n Z X^n | 0 \\rangle(n)$", 'Z')

sx_circuit = QuantumCircuit(1)
sx_circuit.sx(0)

sx_observable = SparsePauliOp.from_list([('Y', 1)])

sx_exact_ev = -1.0

display(sx_circuit.draw("mpl"))
print(sx_observable)
print(sx_exact_ev)

sx_granular_noise_factors = tuple(range(1, 599 + 2, 2))
sx_behavioral_noise_factors = tuple(range(1, 4999, 20))

# with ZneEstimator(
#     circuits=[sx_circuit],
#     observables=[sx_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=sx_granular_noise_factors, noise_amplifier=SingleGateFoldingAmplifier()),
# ) as estimator:
#     sx_granular_a_results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# sx_granular_a_data = {
#     "noise_factors": list(sx_granular_a_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(sx_granular_a_results.metadata[0]["zne"]["values"]),
#     "variances": list(sx_granular_a_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/sx/sx_granular_a.json', 'w') as f:
#     json.dump(sx_granular_a_data, f)

with open('data/sx/sx_granular_a.json', 'r') as f:
    sx_granular_a_data = json.load(f)

# with ZneEstimator(
#     circuits=[sx_circuit],
#     observables=[sx_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=sx_granular_noise_factors, noise_amplifier=SingleGateFoldingAmplifier()),
# ) as estimator:
#     sx_granular_b_results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# sx_granular_b_data = {
#     "noise_factors": list(sx_granular_b_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(sx_granular_b_results.metadata[0]["zne"]["values"]),
#     "variances": list(sx_granular_b_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/sx/sx_granular_b.json', 'w') as f:
#     json.dump(sx_granular_b_data, f)

with open('data/sx/sx_granular_b.json', 'r') as f:
    sx_granular_b_data = json.load(f)

# with ZneEstimator(
#     circuits=[sx_circuit],
#     observables=[sx_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=sx_behavioral_noise_factors, noise_amplifier=SingleGateFoldingAmplifier()),
# ) as estimator:
#     sx_behavioral_a_results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# sx_behavioral_a_data = {
#     "noise_factors": list(sx_behavioral_a_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(sx_behavioral_a_results.metadata[0]["zne"]["values"]),
#     "variances": list(sx_behavioral_a_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/sx/sx_behavioral_a.json', 'w') as f:
#     json.dump(sx_behavioral_a_data, f)

with open('data/sx/sx_behavioral_a.json', 'r') as f:
    sx_behavioral_a_data = json.load(f)

# with ZneEstimator(
#     circuits=[sx_circuit],
#     observables=[sx_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=sx_behavioral_noise_factors, noise_amplifier=SingleGateFoldingAmplifier()),
# ) as estimator:
#     sx_behavioral_b_results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# sx_behavioral_b_data = {
#     "noise_factors": list(sx_behavioral_b_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(sx_behavioral_b_results.metadata[0]["zne"]["values"]),
#     "variances": list(sx_behavioral_b_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/sx/sx_behavioral_b.json', 'w') as f:
#     json.dump(sx_behavioral_b_data, f)

with open('data/sx/sx_behavioral_b.json', 'r') as f:
    sx_behavioral_b_data = json.load(f)

show_experiment([sx_granular_a_data, sx_granular_b_data], sx_exact_ev, "$\sqrt X$ Granular Folding Profile — $\langle 0|{\sqrt X}^nY{\sqrt X}^n|0\\rangle(n)$", 'Y')

show_experiment([sx_behavioral_a_data, sx_behavioral_b_data], sx_exact_ev, "$\sqrt X$ Behavioral Folding Profile — $\langle 0|{\sqrt X}^nY{\sqrt X}^n|0\\rangle(n)$", 'Y')

rz_theta = Parameter('theta')

rz_circuit = QuantumCircuit(1)
rz_circuit.h(0)
rz_circuit.rz(rz_theta, 0)
rz_circuit.h(0)

rz_observable = SparsePauliOp.from_list([('Z', 1)])

rz_theta = np.pi / 4
rz_exact_ev = np.cos(rz_theta)

display(rz_circuit.draw('mpl'))
print(rz_observable)
print(rz_theta)
print(rz_exact_ev)

rz_granular_noise_factors = tuple(range(1, 599 + 2, 2))

# with ZneEstimator(
#     circuits=[rz_circuit],
#     observables=[rz_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=rz_granular_noise_factors, noise_amplifier=SingleGateFoldingAmplifier(gate_name='rz')),
# ) as estimator:
#     rz_granular_a_results = estimator(
#         circuits=[0],
#         observables=[0],
#         parameter_values=[rz_theta]
#     )

# rz_granular_a_data = {
#     "noise_factors": list(rz_granular_a_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(rz_granular_a_results.metadata[0]["zne"]["values"]),
#     "variances": list(rz_granular_a_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/rz/rz_granular_a.json', 'w') as f:
#     json.dump(rz_granular_a_data, f)

with open('data/rz/rz_granular_a.json', 'r') as f:
    rz_granular_a_data = json.load(f)

# with ZneEstimator(
#     circuits=[rz_circuit],
#     observables=[rz_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=rz_granular_noise_factors, noise_amplifier=SingleGateFoldingAmplifier(gate_name='rz')),
# ) as estimator:
#     rz_granular_b_results = estimator(
#         circuits=[0],
#         observables=[0],
#         parameter_values=[rz_theta]
#     )

# rz_granular_b_data = {
#     "noise_factors": list(rz_granular_b_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(rz_granular_b_results.metadata[0]["zne"]["values"]),
#     "variances": list(rz_granular_b_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/rz/rz_granular_b.json', 'w') as f:
#     json.dump(rz_granular_b_data, f)

with open('data/rz/rz_granular_b.json', 'r') as f:
    rz_granular_b_data = json.load(f)

show_experiment([rz_granular_a_data, rz_granular_b_data], rz_exact_ev, "$R_Z$ Folding Profile — $\langle0|H{R_Z(-\\frac\pi4)}^nHZH{R_Z(\\frac\pi4)}^nH|0\\rangle(n)$", 'Z', ev_range=(0.5, 0.9))

cx_circuit = QuantumCircuit(2)
cx_circuit.x(0)
cx_circuit.cx(0, 1)

cx_observable = SparsePauliOp.from_list([('ZI', 1)])

cx_exact_ev = -1.0

display(cx_circuit.draw("mpl"))
print(cx_observable)
print(cx_exact_ev)

cx_granular_noise_factors = tuple(range(1, 599 + 2, 2))
cx_behavioral_noise_factors = tuple(range(1, 5999, 20))

# with ZneEstimator(
#     circuits=[cx_circuit],
#     observables=[cx_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=cx_granular_noise_factors, noise_amplifier=SingleGateFoldingAmplifier(gate_name='cx')),
# ) as estimator:
#     cx_granular_a_results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# cx_granular_a_data = {
#     "noise_factors": list(cx_granular_a_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(cx_granular_a_results.metadata[0]["zne"]["values"]),
#     "variances": list(cx_granular_a_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/cx/cx_granular_a.json', 'w') as f:
#     json.dump(cx_granular_a_data, f)

with open('data/cx/cx_granular_a.json', 'r') as f:
    cx_granular_a_data = json.load(f)

# with ZneEstimator(
#     circuits=[cx_circuit],
#     observables=[cx_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=cx_granular_noise_factors, noise_amplifier=SingleGateFoldingAmplifier(gate_name='cx')),
# ) as estimator:
#     cx_granular_b_results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# cx_granular_b_data = {
#     "noise_factors": list(cx_granular_b_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(cx_granular_b_results.metadata[0]["zne"]["values"]),
#     "variances": list(cx_granular_b_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/cx/cx_granular_b.json', 'w') as f:
#     json.dump(cx_granular_b_data, f)

with open('data/cx/cx_granular_b.json', 'r') as f:
    cx_granular_b_data = json.load(f)

# with ZneEstimator(
#     circuits=[cx_circuit],
#     observables=[cx_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=cx_behavioral_noise_factors, noise_amplifier=SingleGateFoldingAmplifier(gate_name='cx')),
# ) as estimator:
#     cx_behavioral_a_results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# cx_behavioral_a_data = {
#     "noise_factors": list(cx_behavioral_a_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(cx_behavioral_a_results.metadata[0]["zne"]["values"]),
#     "variances": list(cx_behavioral_a_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/cx/cx_behavioral_a.json', 'w') as f:
#     json.dump(cx_behavioral_a_data, f)

with open('data/cx/cx_behavioral_a.json', 'r') as f:
    cx_behavioral_a_data = json.load(f)

# with ZneEstimator(
#     circuits=[cx_circuit],
#     observables=[cx_observable],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=cx_behavioral_noise_factors, noise_amplifier=SingleGateFoldingAmplifier(gate_name='cx')),
# ) as estimator:
#     cx_behavioral_b_results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# cx_behavioral_b_data = {
#     "noise_factors": list(cx_behavioral_b_results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(cx_behavioral_b_results.metadata[0]["zne"]["values"]),
#     "variances": list(cx_behavioral_b_results.metadata[0]["zne"]["variances"]),
# }

# with open('data/cx/cx_behavioral_b.json', 'w') as f:
#     json.dump(cx_behavioral_b_data, f)

with open('data/cx/cx_behavioral_b.json', 'r') as f:
    cx_behavioral_b_data = json.load(f)

show_experiment([cx_granular_a_data, cx_granular_b_data], cx_exact_ev, "$CX$ Granular Folding Profile —  $\langle 00|(XI)({CX})^n(ZI)({CX})^n(XI)|00\\rangle(n)$", 'ZI')

show_experiment([cx_behavioral_a_data, cx_behavioral_b_data], cx_exact_ev, "$CX$ Behavioral Folding Profile —  $\langle 00|(XI)({CX})^n(ZI)({CX})^n(XI)|00\\rangle(n)$", 'ZI')