# 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 extrapolation
from scipy.optimize import curve_fit

# 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.circuit import QuantumCircuit, Parameter
from qiskit.quantum_info import SparsePauliOp
from qiskit.circuit.library import TwoLocal

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

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

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_results(
    measured_data,
    exact_ev,
    experiment_repr,
    observable_repr,
    regime=None,
    regression_model=None,
    regression_guess=None,
    exact_color="#000000",
    measured_color="#dc267f",
    regression_color="#785ef0"
):
    if not regime: regime = len(measured_data["noise_factors"])

    noise_factors = np.array(measured_data["noise_factors"][:regime])
    measured_evs = np.array(measured_data["values"][:regime])
    measured_errors = np.sqrt(np.array(measured_data["variances"][:regime]) / np.array(measured_data["num_shots"][:regime]))

    plt.plot([0, noise_factors[regime-1]], [exact_ev, exact_ev], "--", label=f"Exact", color=exact_color)

    unmitigated_ev = measured_evs[0]
    unmitigated_error = np.abs(exact_ev - unmitigated_ev)
    plt.errorbar([1], measured_evs[0], yerr=measured_errors[0], label=f"Unmitigated ($\Delta\left\langle H \\right\\rangle = {unmitigated_error:.3f}$)", fmt="s", color=measured_color)

    plt.plot(noise_factors, measured_evs, label=f"Measured", color=measured_color)
    plt.fill_between(
        noise_factors,
        measured_evs - measured_errors,
        measured_evs + measured_errors,
        alpha=0.5,
        color=measured_color,
        linewidth=0
    )

    if regression_model:
        popt, pcov = curve_fit(regression_model, measured_data['noise_factors'][:regime], measured_data['values'][:regime], p0=regression_guess)
        goodness_of_fit = np.linalg.norm(np.diag(pcov))
        regression_domain = np.linspace(0, noise_factors[regime-1], 100)
        plt.plot(regression_domain, regression_model(regression_domain, *popt), label=f"Regression", color=regression_color)
        zne_ev = regression_model(0, *popt)
        zne_error = np.abs(exact_ev - zne_ev)
        plt.scatter([0], [zne_ev], label=f"Extrapolated ($\Delta\left\langle H \\right\\rangle = {zne_error:.3f}$)", marker="x", color=regression_color)

    plt.title(experiment_repr)
    plt.xlabel("Noise Factor ($n$)")
    plt.ylabel(f"Expectation Value ($\langle {observable_repr} \\rangle$)")
    plt.legend()
    plt.show()

H2_GNDSTATE_WAVEFN = QuantumCircuit(2).compose(
    TwoLocal(
        rotation_blocks=["ry", "rz"], entanglement_blocks="cz", num_qubits=2
    ).assign_parameters(
        [
            4.104009273873245,
            -6.046872443309065,
            5.277223342332378,
            1.5319564597237907,
            4.893000083985066,
            -3.0921586512224257,
            -0.34939132286472496,
            5.584790805066014,
            -4.161687737508674,
            -3.2391548739165845,
            1.7436977108208878,
            0.17441513089260774,
            2.623695306064223,
            0.03678956439381653,
            2.410055788777767,
            -2.3250358101154913,
        ]
    )
)

H2_HAMILTONIAN = SparsePauliOp.from_list(
    [
        ("II", -1.052373245772859),
        ("IZ", 0.39793742484318045),
        ("ZI", -0.39793742484318045),
        ("ZZ", -0.01128010425623538),
        ("XX", 0.18093119978423156),
    ]
)

H2_GNDSTATE_ENERGY = -1.85727503

NOISE_FACTORS = tuple(range(1, 199 + 2, 2))
len(NOISE_FACTORS)

with Estimator([H2_GNDSTATE_WAVEFN], [H2_HAMILTONIAN]) as estimator:
  simulated_energy = estimator([0], [0]).values[0]
simulated_energy

np.isclose(simulated_energy, H2_GNDSTATE_ENERGY)

# with ZneEstimator(
#     circuits=[H2_GNDSTATE_WAVEFN],
#     observables=[H2_HAMILTONIAN],
#     service=service,
#     options={"backend": backend.name()},
#     zne_strategy=ZNEStrategy(noise_factors=NOISE_FACTORS, noise_amplifier=LocalFoldingAmplifier()),
# ) as estimator:
#     results = estimator(
#         circuits=[0],
#         observables=[0],
#     )

# data = {
#     "noise_factors": list(results.metadata[0]["zne"]["noise_factors"]),
#     "values": list(results.metadata[0]["zne"]["values"]),
#     "variances": list(results.metadata[0]["zne"]["variances"]),
#     "num_shots": [ result["shots"] for result in results.metadata[0]['zne']['remaining_metadata'] ]
# }

# with open('data.json', 'w') as f:
#     json.dump(data, f)

with open('data.json', 'r') as f:
    data = json.load(f)

plot_results(data, H2_GNDSTATE_ENERGY, "$H_2$ Ground State Simulation", "H")

def linear(noise_factor, noiseless_ev, decay_rate):
    return noiseless_ev + decay_rate * noise_factor

for regime in [3, 5, 20]:
    plot_results(data, H2_GNDSTATE_ENERGY, "$H_2$ Ground State Simulation (Linear ZNE)", "H", regime=regime, regression_model=linear)

def damped_sinusoid(noise_factor, noiseless_ev, amplitude, frequency, decay_rate):
    return noiseless_ev + amplitude * (
        1 - np.cos(frequency * noise_factor) * np.exp(-decay_rate * noise_factor)
    )

plot_results(data, H2_GNDSTATE_ENERGY, "$H_2$ Ground State Simulation (Damped Sinusoid ZNE)", "H", regression_model=damped_sinusoid, regression_guess=[data["values"][0], 1, (2*np.pi)/(60*4), 0.01])