from qiskit.circuit import Parameter
from qiskit import QuantumCircuit

theta = Parameter('$\\theta$')

chsh_circuits_no_meas = QuantumCircuit(2)
chsh_circuits_no_meas.h(0)
chsh_circuits_no_meas.cx(0, 1)
chsh_circuits_no_meas.ry(theta, 0)
chsh_circuits_no_meas.draw('mpl')

import numpy as np

number_of_phases = 21
phases = np.linspace(0, 2*np.pi, number_of_phases)
# Phases need to be expressed as list of lists in order to work
individual_phases = [[ph] for ph in phases]

from qiskit_ibm_runtime import QiskitRuntimeService

service = QiskitRuntimeService()
backend = "ibmq_qasm_simulator"  # use the simulator

from qiskit_ibm_runtime import Estimator, Session
from qiskit.quantum_info import SparsePauliOp

ZZ = SparsePauliOp.from_list([("ZZ", 1)])
ZX = SparsePauliOp.from_list([("ZX", 1)])
XZ = SparsePauliOp.from_list([("XZ", 1)])
XX = SparsePauliOp.from_list([("XX", 1)])

ops = [ZZ, ZX, XZ, XX]
chsh_est_sim = []

# Simulator
with Session(service=service, backend=backend):
    estimator = Estimator()
    for op in ops:
        job = estimator.run(
            circuits=[chsh_circuits_no_meas]*len(individual_phases),
            observables=[op]*len(individual_phases),
            parameter_values=individual_phases)
        est_result = job.result()
        chsh_est_sim.append(est_result)

# <CHSH1> = <AB> - <Ab> + <aB> + <ab>
chsh1_est_sim = chsh_est_sim[0].values - chsh_est_sim[1].values + chsh_est_sim[2].values + chsh_est_sim[3].values

# <CHSH2> = <AB> + <Ab> - <aB> + <ab>
chsh2_est_sim = chsh_est_sim[0].values + chsh_est_sim[1].values - chsh_est_sim[2].values + chsh_est_sim[3].values

import matplotlib.pyplot as plt
import matplotlib.ticker as tck

fig, ax = plt.subplots(figsize=(10, 6))
# results from a simulator
ax.plot(phases/np.pi, chsh1_est_sim, 'o-', label='CHSH1 Simulation')
ax.plot(phases/np.pi, chsh2_est_sim, 'o-', label='CHSH2 Simulation')
# classical bound +-2
ax.axhline(y=2, color='r', linestyle='--')
ax.axhline(y=-2, color='r', linestyle='--')
# quantum bound, +-2√2
ax.axhline(y=np.sqrt(2)*2, color='b', linestyle='-.')
ax.axhline(y=-np.sqrt(2)*2, color='b', linestyle='-.')
# set x tick labels to the unit of pi
ax.xaxis.set_major_formatter(tck.FormatStrFormatter('%g $\pi$'))
ax.xaxis.set_major_locator(tck.MultipleLocator(base=0.5))
# set title, labels, and legend
plt.title('Violation of CHSH Inequality')
plt.xlabel('Theta')
plt.ylabel('CHSH witness')
plt.legend()

import qiskit_ibm_runtime
qiskit_ibm_runtime.version.get_version_info()

import qiskit.tools.jupyter

%qiskit_version_table
%qiskit_copyright