from qiskit_ibm_runtime import QiskitRuntimeService, Estimator, Options

# Initialize service
service = QiskitRuntimeService(channel="ibm_quantum")

# Choose a backend to run on (you can change this to another backend that is available)
backend = service.get_backend("ibm_canberra")

from pprint import pprint
from qiskit import QuantumCircuit
from qiskit.circuit import Parameter
import qiskit.quantum_info as qi

# Specify a 1D chain of device qubits to use
initial_layout = [2, 3, 5, 8]
num_qubits = len(initial_layout)

# Parameterized Trotter layer
par_rx = Parameter("a")
par_rzz = Parameter("b")
trotter_layer = QuantumCircuit(num_qubits)
trotter_layer.rx(par_rx, range(num_qubits))
for i in range(num_qubits // 2):
    trotter_layer.rzz(par_rzz, 2 * i, 1 + 2 * i)
for i in range(1, num_qubits // 2 + bool(num_qubits % 2)):
    trotter_layer.rzz(par_rzz, 2 * i - 1, 2 * i)

# Run for 1 to 5 steps
trotter_steps = [1, 2, 3, 4, 5]
trotter_circuit_list = []
for i in trotter_steps:
    trotter_circuit = QuantumCircuit(num_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))

import numpy as np
import qiskit.quantum_info as qi

# Hamiltonian parameters
params = [0.1, -0.2]
params_list = [params] * len(trotter_circuit_list)

# Observable
Zs = np.eye(num_qubits, dtype=bool)
Xs = np.zeros((num_qubits, num_qubits), dtype=bool)
obs = qi.SparsePauliOp(qi.PauliList.from_symplectic(Zs, Xs), 1/num_qubits)
obs_list = [obs]*len(trotter_circuit_list)

# If True run jobs, this can be set to false to avoid re-running jobs when re-running the notebok
# If you are loading saved jobs
run_sim_jobs = True
run_device_jobs = False

# Store all jobs
if run_sim_jobs or run_device_jobs:
    jobs = {}
    jobs_ids = {}

import numpy as np
import scipy.linalg as la
import qiskit.quantum_info as qi
from qiskit_aer.noise import NoiseModel, coherent_unitary_error

backend_simulator = service.get_backend("ibmq_qasm_simulator")

# Noise model with measure noise, depolarizing and thermal relaxation gate noise
noise_model_rodept1 = NoiseModel.from_backend(backend())

# Noise model without measure noise, with depolarizing and thermal relaxation gate noise
noise_model_dept1 = NoiseModel.from_backend(backend(), readout_error=False, gate_error=True, thermal_relaxation=True)

# Noise model with measure noise and with depolarizing gate noise
noise_model_rodep = NoiseModel.from_backend(backend(), readout_error=True, gate_error=True, thermal_relaxation=False)

# Noise model with only measure noise
noise_model_ro = NoiseModel.from_backend(backend(), readout_error=True, gate_error=False, thermal_relaxation=False)

# Noise model with only depolarizing gate noise
noise_model_dep = NoiseModel.from_backend(backend(), readout_error=False, gate_error=True, thermal_relaxation=False)

# Noise model with only coherent CX gate noise
noise_model_coh = NoiseModel.from_backend(backend(), readout_error=False, gate_error=False, thermal_relaxation=False)

# Noise model with measure noise and coherent CX gate noise
noise_model_rocoh = NoiseModel.from_backend(backend(), readout_error=True, gate_error=False, thermal_relaxation=False)

# Construct the coherent CX noise model errors based on the process fidelity so that they have the
# same process fidelity as the depolarizing errors of the depolarizing noise model
for qubits, error in noise_model_dep._local_quantum_errors['cx'].items():
    cos2 = qi.process_fidelity(error)
    theta = np.arccos(np.sqrt(cos2))
    cx_err = coherent_unitary_error(la.expm(-1j * 0.05 * qi.Operator.from_label("ZZ")))
    noise_model_coh.add_quantum_error(cx_err, "cx", qubits)
    noise_model_rocoh.add_quantum_error(cx_err, "cx", qubits)
    
# Labels for noise models
noise_labels = [
    "Ideal",
    "Measure + Depol + Relax",
    "Depol + Relax",
    "Measure + Depol",
    "Measure",
    "Depol",
    "Measure + Coherent",
    "Coherent",
]

# List of labeled noise models
noise_models = [
    None,
    noise_model_rodept1,
    noise_model_dept1,
    noise_model_rodep,
    noise_model_ro,
    noise_model_dep,
    noise_model_rocoh,
    noise_model_coh,
]

# Label for option configuration
label_mit = "Unmitigated"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_sim_jobs:
    # Configure estimator options
    options = Options()
    options.resilience_level = 0
    options.execution.shots = 4 ** 7
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337
    options.simulator.basis_gates = backend.configuration().basis_gates
    options.simulator.coupling_map = backend.configuration().coupling_map
    options.simulator.seed_simulator = 1337

    # Run jobs
    for label, noise_model in zip(noise_labels, noise_models):
        options.simulator.noise_model = noise_model
        estimator = Estimator(backend_simulator, options=options)
        job = estimator.run(trotter_circuit_list, obs_list, params_list)
        jobs[label_mit][label] = job
        print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "Gate Twirling"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_sim_jobs:    
    # Configure estimator options
    options = Options()
    options.resilience_level = 0
    options.twirling.gates = True
    options.execution.shots = 4 ** 7
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337
    options.simulator.basis_gates = backend.configuration().basis_gates
    options.simulator.coupling_map = backend.configuration().coupling_map
    options.simulator.seed_simulator = 1337

    for label, noise_model in zip(noise_labels, noise_models):
        options.simulator.noise_model = noise_model
        estimator = Estimator(backend_simulator, options=options)
        job = estimator.run(trotter_circuit_list, obs_list, params_list)
        jobs[label_mit][label] = job
        print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "TREX wo Gate Twirling"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_sim_jobs:    
    # Configure estimator options
    options = Options()
    options.resilience_level = 1
    options.execution.shots = 4 ** 7
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337
    options.simulator.basis_gates = backend.configuration().basis_gates
    options.simulator.coupling_map = backend.configuration().coupling_map
    options.simulator.seed_simulator = 1337

    for label, noise_model in zip(noise_labels, noise_models):
        options.simulator.noise_model = noise_model
        estimator = Estimator(backend_simulator, options=options)
        job = estimator.run(trotter_circuit_list, obs_list, params_list)
        jobs[label_mit][label] = job
        print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "TREX w Gate Twirling"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_sim_jobs:
    # Configure estimator options
    options = Options()
    options.resilience_level = 1
    options.twirling.gates = True
    options.execution.shots = 4 ** 7
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337
    options.simulator.basis_gates = backend.configuration().basis_gates
    options.simulator.coupling_map = backend.configuration().coupling_map
    options.simulator.seed_simulator = 1337

    for label, noise_model in zip(noise_labels, noise_models):
        options.simulator.noise_model = noise_model
        estimator = Estimator(backend_simulator, options=options)
        job = estimator.run(trotter_circuit_list, obs_list, params_list)
        jobs[label_mit][label] = job
        print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "ZNE w TREX, Gate Twirling"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_sim_jobs:
    # Configure estimator options
    options = Options()
    options.resilience_level = 2
    options.execution.shots = 4 ** 7
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337
    options.simulator.basis_gates = backend.configuration().basis_gates
    options.simulator.coupling_map = backend.configuration().coupling_map
    options.simulator.seed_simulator = 1337

    for label, noise_model in zip(noise_labels, noise_models):
        options.simulator.noise_model = noise_model
        estimator = Estimator(backend_simulator, options=options)
        job = estimator.run(trotter_circuit_list, obs_list, params_list)
        jobs[label_mit][label] = job
        print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "ZNE w TREX wo Gate Twirling"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_sim_jobs:
    # Configure estimator options
    options = Options()
    options.resilience_level = 2
    options.execution.shots = 4 ** 7
    options.twirling.gates = False
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337
    options.simulator.basis_gates = backend.configuration().basis_gates
    options.simulator.coupling_map = backend.configuration().coupling_map
    options.simulator.seed_simulator = 1337

    for label, noise_model in zip(noise_labels, noise_models):
        options.simulator.noise_model = noise_model
        estimator = Estimator(backend_simulator, options=options)
        job = estimator.run(trotter_circuit_list, obs_list, params_list)
        jobs[label_mit][label] = job
        print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "ZNE wo TREX, Gate Twirling"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_sim_jobs:
    # Configure estimator options
    options = Options()
    options.resilience_level = 2
    options.execution.shots = 4 ** 7
    options.twirling.gates = False
    options.twirling.measure = False
    options.resilience.measure_noise_mitigation = False
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337
    options.simulator.basis_gates = backend.configuration().basis_gates
    options.simulator.coupling_map = backend.configuration().coupling_map
    options.simulator.seed_simulator = 1337

    for label, noise_model in zip(noise_labels, noise_models):
        options.simulator.noise_model = noise_model
        estimator = Estimator(backend_simulator, options=options)
        job = estimator.run(trotter_circuit_list, obs_list, params_list)
        jobs[label_mit][label] = job
        print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "PEC"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_sim_jobs:
    # Configure estimator options
    options = Options()
    options.resilience_level = 3
    options.execution.shots = 4 ** 7
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337
    options.simulator.basis_gates = backend.configuration().basis_gates
    options.simulator.coupling_map = backend.configuration().coupling_map
    options.simulator.seed_simulator = 1337

    for label, noise_model in zip(noise_labels, noise_models):
        options.simulator.noise_model = noise_model
        estimator = Estimator(backend_simulator, options=options)
        job = estimator.run(trotter_circuit_list, obs_list, params_list)
        jobs[label_mit][label] = job
        print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "Unmitigated"
label = "Real Device"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_device_jobs:
    
    # Configure estimator options
    options = Options()
    options.resilience_level = 0
    options.execution.shots = 4 ** 7
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337

    estimator = Estimator(backend, options=options)
    job = estimator.run(trotter_circuit_list, obs_list, params_list)
    jobs[label_mit][label] = job
    print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "Gate Twirling"
label = "Real Device"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_device_jobs:    
    # Configure estimator options
    options = Options()
    options.resilience_level = 0
    options.twirling.gates = True
    options.execution.shots = 4 ** 7
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337

    estimator = Estimator(backend, options=options)
    job = estimator.run(trotter_circuit_list, obs_list, params_list)
    jobs[label_mit][label] = job
    print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "TREX wo Gate Twirling"
label = "Real Device"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_device_jobs:    
    # Configure estimator options
    options = Options()
    options.resilience_level = 1
    options.execution.shots = 4 ** 7
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337

    estimator = Estimator(backend, options=options)
    job = estimator.run(trotter_circuit_list, obs_list, params_list)
    jobs[label_mit][label] = job
    print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "TREX w Gate Twirling"
label = "Real Device"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_device_jobs:
    # Configure estimator options
    options = Options()
    options.resilience_level = 1
    options.twirling.gates = True
    options.execution.shots = 4 ** 7
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337

    estimator = Estimator(backend, options=options)
    job = estimator.run(trotter_circuit_list, obs_list, params_list)
    jobs[label_mit][label] = job
    print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "ZNE w TREX, Gate Twirling"
label = "Real Device"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_device_jobs:
    # Configure estimator options
    options = Options()
    options.resilience_level = 2
    options.execution.shots = 4 ** 7
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337

    estimator = Estimator(backend, options=options)
    job = estimator.run(trotter_circuit_list, obs_list, params_list)
    jobs[label_mit][label] = job
    print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "ZNE w TREX wo Gate Twirling"
label = "Real Device"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_device_jobs:
    # Configure estimator options
    options = Options()
    options.resilience_level = 2
    options.execution.shots = 4 ** 7
    options.twirling.gates = False
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337

    estimator = Estimator(backend, options=options)
    job = estimator.run(trotter_circuit_list, obs_list, params_list)
    jobs[label_mit][label] = job
    print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "ZNE wo TREX, Gate Twirling"
label = "Real Device"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_device_jobs:
    # Configure estimator options
    options = Options()
    options.resilience_level = 2
    options.execution.shots = 4 ** 7
    options.twirling.gates = False
    options.twirling.measure = False
    options.resilience.measure_noise_mitigation = False
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337

    estimator = Estimator(backend, options=options)
    job = estimator.run(trotter_circuit_list, obs_list, params_list)
    jobs[label_mit][label] = job
    print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

# Label for option configuration
label_mit = "PEC"
label = "Real Device"

if label_mit not in jobs:
    jobs[label_mit] = {}

if run_device_jobs:
    # Configure estimator options
    options = Options()
    options.resilience_level = 3
    options.execution.shots = 4 ** 7
    options.transpilation.initial_layout = initial_layout
    options.transpilation.seed_transpiler = 1337

    estimator = Estimator(backend, options=options)
    job = estimator.run(trotter_circuit_list, obs_list, params_list)
    jobs[label_mit][label] = job
    print(job.job_id())
    jobs_ids[label_mit] = {label: job.job_id() for label, job in jobs[label_mit].items()}

pprint(jobs_ids)

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

# Configure MPL for plotting
plt.rcParams.update({"text.usetex": True})
plt.rcParams["figure.figsize"] = (6, 4)
mpl.rcParams["figure.dpi"] = 200

load_sim_jobs = True
load_device_jobs = True

from qiskit_ibm_runtime import QiskitRuntimeService

# Re-initialize service if note book has been restarted since jobs ran
service = QiskitRuntimeService(channel="ibm_quantum")

jobs_ids = # Paste job IDs dict here

# Load jobs from job ids
jobs = {}
for label_mit, jids in jobs_ids.items():
    jobs[label_mit] = {}
    for label_noise, jid in jids.items():
        if (label_noise == "Real Device" and load_device_jobs) or (label_noise != "Real Device" and load_sim_jobs):
            jobs[label_mit][label_noise] = service.job(jid)

# Extract expectation value means and standard errors
means_mit = {}
stderrs_mit = {}
for label_mit, jobs_mit in jobs.items():
    means_mit[label_mit] = {}
    stderrs_mit[label_mit] = {}
    for label_noise, job in jobs_mit.items():
        _means = None
        _stderrs = None
        if job.done():
            result = job.result()
            _means = np.array(result.values)
            _stderrs = np.array([metadata["standard_error"] for metadata in result.metadata])
            print(f"LOADED \t| {label_mit}\t| {label_noise}")
        else:
            print(f"{job.status().name}\t| {label_mit}\t| {label_noise}")
        means_mit[label_mit][label_noise] = _means
        stderrs_mit[label_mit][label_noise] = _stderrs

# Prefix for saving figures
fig_prefix = "estimator_trotter_6q"

# Organize mitigation method labels for plotting
labels_mit = [
    'Unmitigated',
    'Gate Twirling',
    'TREX w Gate Twirling',
    'TREX wo Gate Twirling',
    'ZNE wo TREX, Gate Twirling',
    'ZNE w TREX wo Gate Twirling',
    'ZNE w TREX, Gate Twirling',
    'PEC',
]

# Organize noise model labels for plotting
labels_noise = [
    'Ideal',
    'Real Device',
    'Measure',
    'Depol',
    'Depol + Relax',
    'Coherent',
    'Measure + Depol',
    'Measure + Depol + Relax',
    'Measure + Coherent',
]

# Generate subplots 
fig, axs = plt.subplots(3, 3, sharex=True, sharey=True, figsize=(10, 10))

# Flatten the array of axes for easy iteration
axs = axs.ravel()

# Add error bar plots to each subplot
for label_noise in labels_noise:
    if label_noise == "Real Device":
        axs[0].plot([5], [0], 's-', c="red", label=label_noise)
    elif label_noise == "Ideal":
        continue
    else:
        axs[0].plot([5], [0], 'o--', label=label_noise)
axs[0].legend()
for ax, label_mit in zip(axs[1:], labels_mit):
    means = means_mit[label_mit]
    stderrs = stderrs_mit[label_mit]
    for label_noise in labels_noise:
        if label_noise != "Ideal" and means[label_noise] is not None and means["Ideal"] is not None:
            ys = means["Ideal"] - means[label_noise]
            yerrs = stderrs["Ideal"] + stderrs[label_noise]
            xs = 1 + np.arange(len(ys))
            if label_noise == "Real Device":
                ax.errorbar(xs, ys, yerrs, fmt = 's-', c="red", capsize=4, label=label_noise)
            else:
                ax.errorbar(xs, ys, yerrs, fmt = 'o--', capsize=4, label=label_noise)
    ax.plot(xs, np.zeros_like(xs), ":", c="black")
    ax.set_title(label_mit)
    ax.set_xticks(xs)

plt.tight_layout()
plt.savefig(f"{fig_prefix}_mit_vs_noise.pdf")
plt.show()

# Generate subplots 
fig, axs = plt.subplots(3, 3, sharex=True, sharey=True, figsize=(10, 10))

# Flatten the array of axes for easy iteration
axs = axs.ravel()

# Add error bar plots to each subplot
for ax, label_noise in zip(axs, labels_noise):
    for label_mit in labels_mit:
        means = means_mit[label_mit]
        stderrs = stderrs_mit[label_mit]
        if means[label_noise] is not None and means["Ideal"] is not None:
            ys = means["Ideal"] - means[label_noise]
            yerrs = stderrs["Ideal"] + stderrs[label_noise]
            xs = 1 + np.arange(len(ys))
            if label_mit == "Unmitigated":
                ax.errorbar(xs, ys, yerrs, fmt = 's-', capsize=4, label=label_mit)
            else:
                ax.errorbar(xs, ys, yerrs, fmt = 'o--', capsize=4, label=label_mit)
    ax.plot(xs, np.zeros_like(xs), ":", c="black")
    if label_noise == "Real Device":
        title = label_noise
    else:
        title = f"{label_noise} Sim"
    ax.set_title(title)
    ax.set_xticks(xs)
    if label_noise == "Ideal":
        ax.legend()
plt.tight_layout()
plt.savefig(f"{fig_prefix}_noise_vs_mit.pdf")
plt.show()