import datetime
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
plt.rcParams.update({"text.usetex": True})
plt.rcParams["figure.figsize"] = (6,4)
mpl.rcParams["figure.dpi"] = 200

from qiskit_ibm_runtime import Estimator, Session, QiskitRuntimeService, Options
from qiskit.quantum_info import SparsePauliOp
from qiskit import QuantumCircuit

service = QiskitRuntimeService()

backend_simulator = "backend_simulator"
backend = "ibmq_montreal"

qubits = 4
trotter_layer = QuantumCircuit(qubits)
trotter_layer.rx(0.1, range(qubits))
trotter_layer.cx(0, 1)
trotter_layer.cx(2, 3)
trotter_layer.rz(-0.2, [1, 3])
trotter_layer.cx(0, 1)
trotter_layer.cx(2, 3)
trotter_layer.cx(1, 2)
trotter_layer.rz(-0.2, 2)
trotter_layer.cx(1, 2)

num_steps = 6
trotter_circuit_list = []
for i in range(1, num_steps):
    trotter_circuit = QuantumCircuit(qubits)
    for _ in range(i):
        trotter_circuit = trotter_circuit.compose(trotter_layer)
    trotter_circuit_list.append(trotter_circuit)
    print(f'Trotter circuit with {i} Trotter steps`)
    display(trotter_circuit.draw(fold=-1))

obs = SparsePauliOp("Z"*qubits)
obs_list = [obs]*len(trotter_circuit_list)


options = Options()
options.execution.shots = 1000
options.optimization_level = 0 # No optimization
options.resilience_level = 0 # No mitigation

with Session(service=service, backend=backend_simulator) as session:
    estimator_sim = Estimator(session=session, options=options)
    job_sim = estimator_sim.run(circuits=trotter_circuit_list, observables=obs_list)
    print('job id:', job_sim.job_id)
    print(job_sim.result())

expvals_ideal = job_sim.result().values
expvals_ideal_variance = [metadata['variance']/metadata['shots'] for metadata in job_sim.result().metadata]
std_error_ideal = np.sqrt(expvals_ideal_variance)

options = Options()
options.execution.shots = 1000
options.optimization_level = 0 # No optimization
options.resilience_level = 0 # No error mitigation

with Session(service=service, backend=backend) as session:
    estimator = Estimator(session=session, options=options)
    job = estimator.run(circuits=trotter_circuit_list, observables=obs_list)
    print('job id:', job.job_id)
    print(job.result())

expvals_unmit = job.result().values
expvals_unmit_variance = [metadata['variance']/metadata['shots'] for metadata in job.result().metadata]
std_error_unmit = np.sqrt(expvals_unmit_variance)

options = Options()
options.execution.shots = 1000
options.optimization_level = 3 # Dynamical decoupling
options.resilience_level = 0 # No error mitigation

with Session(service=service, backend=backend) as session:
    estimator = Estimator(session=session, options=options)
    job_dd = estimator.run(circuits=trotter_circuit_list, observables=obs_list)
    print('job id:', job_dd.job_id)
    print(job_dd.result())

expvals_unmit_dd = job_dd.result().values
expvals_unmit_dd_variance = [metadata['variance']/metadata['shots'] for metadata in job_dd.result().metadata]
std_error_dd = np.sqrt(expvals_unmit_dd_variance)

plt.title('Trotter circuits expectation value')
plt.errorbar(range(1, num_steps), expvals_ideal, std_error_ideal, fmt = 'o', linestyle = '--', capsize=4, c='red', label='Ideal')
plt.errorbar(range(1, num_steps), expvals_unmit, std_error_unmit, fmt = 'o', linestyle = '-', capsize=4, c='green', label='No mitigation')
plt.errorbar(range(1, num_steps), expvals_unmit_dd, std_error_dd, fmt = 'o', linestyle = '-', capsize=4, c='blue', label='Dynamical decoupling')

plt.ylabel(f"$\langle ZZZZ \\rangle$")
plt.xlabel('No. Trotter Steps')
plt.xticks([1, 2, 3, 4, 5])
plt.legend()
plt.show()

options = Options()
options.resilience_level = 1 # T-REx
options.optimization_level = 0 # No optimization
options.execution.shots = 1000

with Session(service=service, backend=backend) as session:
    estimator = Estimator(session=session, options=options)
    job_trex = estimator.run(circuits=trotter_circuit_list, observables=obs_list)
    print('job id:', job_trex.job_id)
    print(job_trex.result())

expvals_unmit_trex = job_trex.result().values
expvals_unmit_trex_variance = [metadata['variance']/metadata['shots'] for metadata in job_trex.result().metadata]
std_error_trex = np.sqrt(expvals_unmit_trex_variance)

plt.title('Trotter circuits expectation value')
plt.errorbar(range(1, num_steps), expvals_ideal, std_error_ideal, fmt = 'o', linestyle = '--', capsize=4, c='red', label='Ideal')
plt.errorbar(range(1, num_steps), expvals_unmit, std_error_unmit, fmt = 'o', linestyle = '-', capsize=4, c='green', label='No mitigation')
plt.errorbar(range(1, num_steps), expvals_unmit_trex, std_error_trex, fmt = 'o', linestyle = '-', capsize=4, c='violet', label='T-REx')

plt.ylabel(f"$\langle ZZZZ \\rangle$")
plt.xlabel('No. Trotter Steps')
plt.xticks([1, 2, 3, 4, 5])
plt.legend()
plt.show()

options = Options()
options.execution.shots = 1000
options.optimization_level = 0 # No optimization
options.resilience_level = 2 # ZNE

with Session(service=service, backend=backend) as session:
    estimator = Estimator(session=session, options=options)
    job_zne = estimator.run(circuits=trotter_circuit_list, observables=obs_list)
    print('job id:', job_zne.job_id)
    print(job_zne.result())

expvals_unmit_zne = job_zne.result().values
# Standard error: coming soon!


plt.title('Trotter circuits expectation value')
plt.errorbar(range(1, num_steps), expvals_ideal, std_error_ideal, fmt = 'o', linestyle = '--', capsize=4, c='red', label='Ideal')
plt.errorbar(range(1, num_steps), expvals_unmit, std_error_unmit, fmt = 'o', linestyle = '-', capsize=4, c='green', label='No mitigation')
plt.errorbar(range(1, num_steps), expvals_unmit_zne, [0]*(num_steps-1), fmt = 'o', linestyle = '-', capsize=4, c='cyan', label='ZNE')

plt.xlabel('No. Trotter Steps')
plt.ylabel(f"$\langle ZZZZ \\rangle$")
plt.xticks([1, 2, 3, 4, 5])
plt.legend()
plt.show()

def interim_results_callback(job_id, result):
    now = datetime.datetime.now()
    print(now, "*** Callback ***", result, "\n")

options = Options()
options.optimization_level = 0 # No optimization
options.execution.shots = 100
options.resilience_level = 3 # PEC
options.environment.callback = interim_results_callback

with Session(service=service, backend=backend) as session:
    estimator_pec = Estimator(session=session, options=options)
    job_pec = estimator_pec.run(circuits=trotter_circuit_list, observables=obs_list)
    print('job id:', job_pec.job_id)

expvals_pec = job_pec.result().values
std_error_pec = [metadata['standard_error'] for metadata in job_pec.result().metadata]


plt.title('Trotter circuits expectation value')
plt.errorbar(range(1, num_steps), expvals_ideal, std_error_ideal, fmt = 'o', linestyle = '--', capsize=4, c='red', label='Ideal')
plt.errorbar(range(1, num_steps), expvals_unmit, std_error_unmit, fmt = 'o', linestyle = '-', capsize=4, c='green', label='No mitigation')
plt.errorbar(range(1, num_steps), expvals_pec, std_error_pec, fmt = 'd', linestyle = '-', capsize=4, c='orange', label='PEC')

plt.ylabel(f"$\langle ZZZZ \\rangle$")
plt.xlabel('No. Trotter Steps')
plt.xticks([1, 2, 3, 4, 5])
plt.legend()
plt.show()

print(job_pec.result())

pec_metadata = job_pec.result().metadata
fig, ax = plt.subplots()
fig.subplots_adjust(right=0.75)

twin1 = ax.twinx()
twin2 = ax.twinx()
twin3 = ax.twinx()

twin2.spines.right.set_position(("axes", 1.2))
twin3.spines.right.set_position(("axes", 1.4))

p1, = ax.plot(range(1, num_steps), [m["total_mitigated_layers"] for m in pec_metadata] , "b-", label="Total mitigated layers")
p2, = twin1.plot(range(1, num_steps), [m["sampling_overhead"] for m in pec_metadata], "r-", label="Sampling overhead")
p3, = twin2.plot(range(1, num_steps), [m["samples"] for m in pec_metadata], "g-", label="Samples")
p4, = twin3.plot(range(1, num_steps), [m["shots"] for m in pec_metadata], "c-", label="Shots")

ax.set_ylim(0, 20)
twin1.set_ylim(0, 2.8)
twin2.set_ylim(0, 300)
twin3.set_ylim(0, 35000)

ax.set_xlabel("No. Trotter Steps")
ax.set_ylabel("Total mitigated layers")
twin1.set_ylabel("Sampling overhead")
twin2.set_ylabel("Samples")
twin3.set_ylabel("Shots")

ax.yaxis.label.set_color(p1.get_color())
twin1.yaxis.label.set_color(p2.get_color())
twin2.yaxis.label.set_color(p3.get_color())
twin3.yaxis.label.set_color(p4.get_color())

tkw = dict(size=4, width=1.5)
ax.tick_params(axis='y', colors=p1.get_color(), **tkw)
twin1.tick_params(axis='y', colors=p2.get_color(), **tkw)
twin2.tick_params(axis='y', colors=p3.get_color(), **tkw)
twin3.tick_params(axis='y', colors=p4.get_color(), **tkw)
plt.xticks([1, 2, 3, 4, 5])

ax.legend(handles=[p1, p2, p3, p4])

plt.title('PEC metadata')
plt.show()

from matplotlib.pyplot import figure
plt.errorbar(range(1, num_steps), expvals_ideal, std_error_ideal, fmt = 'o', linestyle = '--', capsize=4, c='red', label='Ideal')
plt.errorbar(range(1, num_steps), expvals_unmit, std_error_unmit, fmt = 'o', linestyle = '-', capsize=4, c='green', label='No mitigation')
plt.errorbar(range(1, num_steps), expvals_unmit_trex, std_error_trex, fmt = 'o', linestyle = '-', capsize=4, c='violet', label='T-REx')
plt.errorbar(range(1, num_steps), expvals_unmit_zne, [0]*(num_steps-1), fmt = 'o', linestyle = '-', capsize=4, c='cyan', label='ZNE')
plt.errorbar(range(1, num_steps), expvals_pec, std_error_pec, fmt = 'd', linestyle = '-', capsize=4, c='orange', label='PEC')

plt.title('Trotter circuits expectation value')
plt.ylabel(f"$\langle ZZZZ \\rangle$")
plt.xlabel('No. Trotter Steps')
plt.xticks([1, 2, 3, 4, 5])
plt.legend()
plt.show()

options = Options()
options.execution.shots = 1000
options.optimization_level = 0 # no optimization
options.resilience_level = 2 # ZNE
options.resilience.noise_factors = [1, 2, 3, 4]
options.resilience.noise_amplifier = "LocalFoldingAmplifier"
options.resilience.extrapolator = "QuadraticExtrapolator"

with Session(service=service, backend='ibmq_montreal') as session:
    estimator = Estimator(session=session, options=options)
    job_zne_options = estimator.run(circuits=trotter_circuit_list, observables=obs_list)
    print('job id:', job_zne_options.job_id)
    print(job_zne_options.result())

from qiskit.tools import jupyter
%qiskit_version_table
%qiskit_copyright