# load necessary Runtime libraries
from qiskit_ibm_runtime import QiskitRuntimeService, Sampler, Session

backend = "ibmq_qasm_simulator"  # use the simulator

from qiskit.circuit import Parameter
from qiskit.opflow import I, X, Z

mu = Parameter('$\\mu$')
ham_pauli = mu * X

cc = Parameter('$c$')
ww = Parameter('$\\omega$')

ham_res = -(1/2)*ww*(I^Z) + cc*(X^X) + (ham_pauli^I)

tt = Parameter('$t$')
U_ham = (tt*ham_res).exp_i()

from qiskit import transpile
from qiskit.circuit import ClassicalRegister
from qiskit.opflow import PauliTrotterEvolution, Suzuki
import numpy as np

num_trot_steps = 5
total_time = 10
cr = ClassicalRegister(1, 'c')

spec_op = PauliTrotterEvolution(trotter_mode=Suzuki(order=2, reps=num_trot_steps)).convert(U_ham)
spec_circ = spec_op.to_circuit()
spec_circ_t = transpile(spec_circ, basis_gates=['sx', 'rz', 'cx'])
spec_circ_t.add_register(cr)
spec_circ_t.measure(0, cr[0])

spec_circ_t.draw('mpl')

# fixed Parameters
fixed_params = {
    cc: 0.3,
    mu: 0.7,
    tt: total_time
}
# Parameter value for single circuit
param_keys = list(spec_circ_t.parameters)

# run through all the ww values to create a List of Lists of Parameter value
num_pts = 101
wvals = np.linspace(-2, 2, num_pts)
param_vals = []
for wval in wvals:
    all_params = {**fixed_params, **{ww: wval}}
    param_vals.append([all_params[key] for key in param_keys])

with Session(backend=backend):
    sampler = Sampler()
    job = sampler.run(
        circuits=[spec_circ_t]*num_pts,
        parameter_values=param_vals,
        shots=1e5
    )
    result = job.result()


Zexps = []
for dist in result.quasi_dists:
    if 1 in dist:
        Zexps.append(1 - 2*dist[1])
    else:
        Zexps.append(1)

from qiskit.opflow import PauliExpectation, Zero

param_bind = {
    cc: 0.3,
    mu: 0.7,
    tt: total_time
}

init_state = Zero^2
obsv = I^Z
Zexp_exact = (U_ham @ init_state).adjoint() @ obsv @ (U_ham @ init_state)

diag_meas_op = PauliExpectation().convert(Zexp_exact)
Zexact_values = []
for w_set in wvals:
    param_bind[ww] = w_set
    Zexact_values.append(np.real(diag_meas_op.bind_parameters(param_bind).eval()))

import matplotlib.pyplot as plt
plt.style.use('dark_background')

fig, ax = plt.subplots(dpi=100)
ax.plot([-param_bind[mu], -param_bind[mu]], [0, 1], ls='--', color='purple')
ax.plot([param_bind[mu], param_bind[mu]], [0, 1], ls='--', color='purple')
ax.plot(wvals, Zexact_values, label='Exact')
ax.plot(wvals, Zexps, label=f"{backend}")
ax.set_xlabel(r'$\omega$ (arb)')
ax.set_ylabel(r'$\langle Z \rangle$ Expectation')
ax.legend()

import qiskit_ibm_runtime
qiskit_ibm_runtime.version.get_version_info()

from qiskit.tools.jupyter import *
%qiskit_version_table