import networkx as nx
graph = nx.Graph()
graph.add_nodes_from([0, 1, 2, 3, 4])
for i, j in [(0, 1), (0,4), (1, 2), (1, 4) , (2, 3) , (3, 4)]:
    graph.add_edge(i, j)

from qat.generators import MaxCutGenerator
from qat.qpus import SQAQPU
max_cut_application = MaxCutGenerator(job_type="annealing") | SQAQPU()
combinatorial_result = max_cut_application.execute(graph)

from qat.generators import MaxCutGenerator
from qat.plugins import ScipyMinimizePlugin
from qat.qpus import get_default_qpu
scipy_args = dict(method="COBYLA", tol=1e-5, options={"maxiter": 200})
max_cut_application = MaxCutGenerator(job_type="qaoa") | (ScipyMinimizePlugin(**scipy_args) | get_default_qpu())
combinatorial_result = max_cut_application.execute(graph)

from qlmaas.plugins import Nnizer
from qat.core import Batch
nnizer_atos = Nnizer(method="atos")
nnized_batch_atos = nnizer_atos.compile(Batch(jobs=[job]), my_hardware)
nnized_ansatz_circuit_atos = nnized_batch_atos.jobs[0].circuit
%qatdisplay nnized_ansatz_circuit_atos --svg

from qat.hardware import make_depolarizing_hardware_model
hw_model = make_depolarizing_hardware_model(eps1=0.01, eps2=0.05)
from qlmaas.qpus import NoisyQProc
noisy_qpu = NoisyQProc(hardware_model=hw_model)
job_2 = circuit_1.to_job()
result_2 = noisy_qpu.submit(job_2)

 
from qat.core import Observable, Term
nbqbits = 4
pauli_terms = [Term(1, op, [0, 1, 2, 3]) for op in ["ZZZZ"]]
obs = Observable(nbqbits, pauli_terms=pauli_terms)
job = bd_circ.to_job(observable=obs)
perfect_qpu = LinAlg()
perfect_res = perfect_qpu.submit(job)
E_noisefree = perfect_res.value
print("E without noise=", E_noisefree)

eps1 = 0.016
eps2 = 0.085
hw_model = make_depolarizing_hardware_model(eps1=eps1, eps2=eps2, depol_type="randomizing", correl_type="multi_qubit")
noisy_qpu = NoisyQProc(hardware_model=hw_model, sim_method="deterministic-vectorized")
noisy_res = noisy_qpu.submit(job)
E_noisy = noisy_res.value
print("E noisy=", E_noisy, "(error: %s %%)" % (abs(E_noisefree - E_noisy) / abs(E_noisefree) * 100))

ZNE_stack = ZeroNoiseExtrapolator(n_ins=1) | noisy_qpu
ZNE_res = ZNE_stack.submit(job)
E_mitigated_linear = ZNE_res.value
print(
    "E mitigated with linear ansatz:",
    E_mitigated_linear,
    "(error: %s %%)" % (abs(E_noisefree - E_mitigated_linear) / abs(E_noisefree) * 100),
)
fit_parameters = eval(ZNE_res.meta_data["ZNE_fit_parameters"])
a = fit_parameters["a"]
b = fit_parameters["b"]
linear_fit = [a * x + b for x in range(-1, 100)]

ZNE_stack = ZeroNoiseExtrapolator(n_ins=1, extrap_method="exponential") | noisy_qpu
ZNE_res = ZNE_stack.submit(job)
E_mitigated_exponential = ZNE_res.value
print(
    "E mitigated with exponential ansatz:",
    E_mitigated_exponential,
    "(error: %s %%)" % (abs(E_noisefree - E_mitigated_exponential) / abs(E_noisefree) * 100),
)
fit_parameters = eval(ZNE_res.meta_data["ZNE_fit_parameters"])
a = fit_parameters["a"]
b = fit_parameters["b"]
exp_fit = [np.exp(a * x + b) + obs.constant_coeff for x in range(-1, 100)]
# be careful with the sign before the exponential, it depends on the job!

plt.plot(energies, marker="o", label="noisy QPU energies")
plt.plot(range(-1, 100), exp_fit, ls="dashed", label="fit")
plt.axhline(obs.constant_coeff, ls="dashed", color="sienna", label="$\mathrm{Tr}(H)/2^n$")
plt.axhline(E_noisefree, ls="dashed", color="green", label="perfect energy")
plt.axhline(E_mitigated_exponential, ls="dashed", color="purple", label="ZNE energy")
plt.grid()
plt.legend()
plt.ylim(ymin=obs.constant_coeff - 0.01, ymax=perfect_res.value + 0.01)
plt.xlabel("$n_{\mathrm{pairs}}$")
plt.ylabel("energy")