from qiskit import QuantumCircuit
def qc_state_prep():
    qc = QuantumCircuit(2, name = 'state_prep'); 
    qc.h(0)
    qc.cx(0, 1)
    qc.z(0)
    qc.x(1)
    return qc
qc_state_prep_ = qc_state_prep(); qc_state_prep_.draw(output = 'mpl')

from qiskit import IBMQ, Aer, execute
nshots = 8192
provider = IBMQ.load_account()
device = provider.get_backend('ibmq_quito')
simulator = Aer.get_backend('qasm_simulator')
from qiskit.tools.monitor import job_monitor
from qiskit.ignis.verification.tomography import state_tomography_circuits, StateTomographyFitter
from qiskit.visualization import plot_state_city

qstc = state_tomography_circuits(qc_state_prep_, [0, 1])
job = execute(qstc, backend = simulator, shots = nshots)
qstf = StateTomographyFitter(job.result(), qstc)
rho = qstf.fit(method = 'lstsq')
plot_state_city(rho, title = r'$|\Psi_{-}^{sim}\rangle$')

import math

def qc_A1B1():
    qc = QuantumCircuit(2, name = 'A1B1'); qc.z(0); qc.u(math.pi/4, 0, math.pi, [1])
    return qc
qc_A1B1_ = qc_A1B1(); qc_A1B1_.draw(output = 'mpl')

qc = QuantumCircuit(2, 2)
qc_state_prep_ = qc_state_prep(); qc.append(qc_state_prep_, [0, 1])
qc.barrier()
qc_A1B1_ = qc_A1B1(); qc.append(qc_A1B1_, [0, 1]); qc.measure([0, 1], [0, 1])
qc.draw(output = 'mpl')

jobS = execute(qc, backend = simulator, shots = nshots)
resultS = jobS.result().get_counts(); print(resultS , ',', resultS['11'])

def AB_avg(result, nshots):
    avg = 0
    if '00' in result:
        avg += result['00']
    if '01' in result:
        avg -= result['01']
    if '10' in result:
        avg -= result['10']
    if '11' in result:
        avg += result['11']
    return avg/nshots

A1B1_avg = AB_avg(resultS, nshots); print(A1B1_avg)

jobE = execute(qc, backend = device, shots = nshots); job_monitor(jobE)
A1B1_avgE = AB_avg(jobE.result().get_counts(), nshots)

print('<A1B1>_teo = ', -1/math.sqrt(2), ', <A1B1>_sim = ', A1B1_avg, ', <A1B1>_exp = ', A1B1_avgE)

def qc_A1B2():
    qc = QuantumCircuit(2, name = 'A1B2'); qc.z(0); qc.u(-math.pi/4, 0, math.pi, [1])
    return qc
qc_A1B2_ = qc_A1B2(); qc_A1B2_.draw(output = 'mpl')

qc = QuantumCircuit(2, 2)
qc_state_prep_ = qc_state_prep(); qc.append(qc_state_prep_, [0, 1])
qc_A1B2_ = qc_A1B2(); qc.append(qc_A1B2_, [0, 1]); qc.measure([0, 1], [0, 1])
jobS = execute(qc, backend = simulator, shots = nshots)
resultS = jobS.result().get_counts()
A1B2_avg = AB_avg(resultS, nshots)
jobE = execute(qc, backend = device, shots = nshots); job_monitor(jobE)
A1B2_avgE = AB_avg(jobE.result().get_counts(), nshots)
print('<A1B2>_teo = ', -1/math.sqrt(2), ', <A1B2>_sim = ', A1B2_avg, ', <A1B2>_exp = ', A1B2_avgE)

def qc_A2B1():
    qc = QuantumCircuit(2, name = 'A2B1'); qc.h(0); qc.u(math.pi/4, 0, math.pi, [1])
    return qc
qc_A2B1_ = qc_A2B1(); qc_A2B1_.draw(output = 'mpl')

qc = QuantumCircuit(2, 2)
qc_state_prep_ = qc_state_prep(); qc.append(qc_state_prep_, [0, 1])
qc_A2B1_ = qc_A2B1(); qc.append(qc_A2B1_, [0, 1]); qc.measure([0, 1], [0, 1])
jobS = execute(qc, backend = simulator, shots = nshots)
resultS = jobS.result().get_counts()
A2B1_avg = AB_avg(resultS, nshots)
jobE = execute(qc, backend = device, shots = nshots); job_monitor(jobE)
A2B1_avgE = AB_avg(jobE.result().get_counts(), nshots)
print('<A2B1>_teo = ', -1/math.sqrt(2), ', <A2B1>_sim = ', A2B1_avg, ', <A2B1>_exp = ', A2B1_avgE)

def qc_A2B2():
    qc = QuantumCircuit(2, name = 'A2B2'); qc.h(0); qc.u(-math.pi/4, 0, math.pi, [1])
    return qc
qc_A2B2_ = qc_A2B2(); qc_A2B2_.draw(output = 'mpl')

qc = QuantumCircuit(2, 2)
qc_state_prep_ = qc_state_prep(); qc.append(qc_state_prep_, [0, 1])
qc_A2B2_ = qc_A2B2(); qc.append(qc_A2B2_, [0, 1]); qc.measure([0, 1], [0, 1])
jobS = execute(qc, backend = simulator, shots = nshots)
resultS = jobS.result().get_counts()
A2B2_avg = AB_avg(resultS, nshots)
jobE = execute(qc, backend = device, shots = nshots); job_monitor(jobE)
A2B2_avgE = AB_avg(jobE.result().get_counts(), nshots)
print('<A2B2>_teo = ', 1/math.sqrt(2), ', <A2B2>_sim = ', A2B2_avg, ', <A2B2>_exp = ', A2B2_avgE)

O_avg = A1B1_avg + A1B2_avg + A2B1_avg - A2B2_avg
O_avgE = A1B1_avgE + A1B2_avgE + A2B1_avgE - A2B2_avgE
print('<O>_teo = ', -2*math.sqrt(2), ', <O>_sim = ', O_avg, ', <O>_exp = ', O_avgE)

import numpy as np
from matplotlib import pyplot as plt
import matplotlib

def qc_state_prep(th):
    qc = QuantumCircuit(2, name = 'Psi')
    qc.u(th, 0, 0, [0]); qc.cx([0], [1]); qc.z([0]); qc.x([1])
    return qc
qc_state_prep_ = qc_state_prep(math.pi); qc_state_prep_.draw(output = 'mpl')

p = np.arange(0, math.pi/2+0.1, 0.1); Oavg = -math.sqrt(2)*(1+np.sin(p))
d = p.shape[0]; y = -2*np.ones(d); Oavg_sim = np.zeros(d); Oavg_exp = np.zeros(d)
for j in range(0, d):
    # A1B1
    qc = QuantumCircuit(2, 2)
    qc_state_prep_ = qc_state_prep(p[j]); qc.append(qc_state_prep_, [0, 1])
    qc_A1B1_ = qc_A1B1(); qc.append(qc_A1B1_, [0, 1]); qc.measure([0, 1], [0, 1])
    jobS = execute(qc, backend = simulator, shots = nshots)
    resultS = jobS.result().get_counts(); A1B1_avg = AB_avg(resultS, nshots)
    jobE = execute(qc, backend = device, shots = nshots); job_monitor(jobE)
    A1B1_avgE = AB_avg(jobE.result().get_counts(), nshots)
    # A1B2
    qc = QuantumCircuit(2, 2)
    qc_state_prep_ = qc_state_prep(p[j]); qc.append(qc_state_prep_, [0, 1])
    qc_A1B2_ = qc_A1B2(); qc.append(qc_A1B2_, [0, 1]); qc.measure([0, 1], [0, 1])
    jobS = execute(qc, backend = simulator, shots = nshots)
    resultS = jobS.result().get_counts(); A1B2_avg = AB_avg(resultS, nshots)
    jobE = execute(qc, backend = device, shots = nshots); job_monitor(jobE)
    A1B2_avgE = AB_avg(jobE.result().get_counts(), nshots)
    # A2B1
    qc = QuantumCircuit(2, 2)
    qc_state_prep_ = qc_state_prep(p[j]); qc.append(qc_state_prep_, [0, 1])
    qc_A2B1_ = qc_A2B1(); qc.append(qc_A2B1_, [0, 1]); qc.measure([0, 1], [0, 1])
    jobS = execute(qc, backend = simulator, shots = nshots)
    resultS = jobS.result().get_counts(); A2B1_avg = AB_avg(resultS, nshots)
    jobE = execute(qc, backend = device, shots = nshots); job_monitor(jobE)
    A2B1_avgE = AB_avg(jobE.result().get_counts(), nshots)
    # A2B2
    qc = QuantumCircuit(2, 2)
    qc_state_prep_ = qc_state_prep(p[j]); qc.append(qc_state_prep_, [0, 1])
    qc_A2B2_ = qc_A2B2(); qc.append(qc_A2B2_, [0, 1]); qc.measure([0, 1], [0, 1])
    jobS = execute(qc, backend = simulator, shots = nshots)
    resultS = jobS.result().get_counts(); A2B2_avg = AB_avg(resultS, nshots)
    jobE = execute(qc, backend = device, shots = nshots); job_monitor(jobE)
    A2B2_avgE = AB_avg(jobE.result().get_counts(), nshots)
    # O
    Oavg_sim[j] = A1B1_avg + A1B2_avg + A2B1_avg - A2B2_avg
    Oavg_exp[j] = A1B1_avgE + A1B2_avgE + A2B1_avgE - A2B2_avgE

matplotlib.rcParams.update({'font.size':12})
plt.figure(figsize = (6,4), dpi = 100)
plt.plot(p, y, label = 'LHV limit')
plt.plot(p, Oavg, label = r'$\langle O\rangle_{\Psi}^{teo}$')
plt.plot(p, Oavg_sim, '*', label = r'$\langle O\rangle_{\Psi}^{sim}$')
plt.plot(p, Oavg_exp, 'o', label = r'$\langle O\rangle_{\Psi}^{exp}$')
plt.xlabel(r'$\theta$'); plt.legend(); plt.show()

def qc_A1B1(th):
    qc = QuantumCircuit(2, name = 'A1B1'); qc.z(0); qc.u(th, 0, math.pi, [1])
    return qc
def qc_A1B2():
    qc = QuantumCircuit(2, name = 'A1B2'); qc.z(0); qc.u(-math.pi/4, 0, math.pi, [1])
    return qc
def qc_A2B1(th):
    qc = QuantumCircuit(2, name = 'A2B1'); qc.h(0); qc.u(th, 0, math.pi, [1])
    return qc
def qc_A2B2():
    qc = QuantumCircuit(2, name = 'A2B2'); qc.h(0); qc.u(-math.pi/4, 0, math.pi, [1])
    return qc
def qc_state_prep():
    qc = QuantumCircuit(2, name = 'state_prep'); qc.h(0); qc.cx(0, 1); qc.z(0); qc.x(1)
    return qc

th = np.arange(-math.pi/4, math.pi/4+0.1, 0.1); Oavg = -(math.sqrt(2)+np.cos(th)+np.sin(th))
d = p.shape[0]; y = -2*np.ones(d); Oavg_sim = np.zeros(d); Oavg_exp = np.zeros(d)
for j in range(0, d):
    # A1B1
    qc = QuantumCircuit(2, 2)
    qc_state_prep_ = qc_state_prep(); qc.append(qc_state_prep_, [0, 1])
    qc_A1B1_ = qc_A1B1(th[j]); qc.append(qc_A1B1_, [0, 1]); qc.measure([0, 1], [0, 1])
    jobS = execute(qc, backend = simulator, shots = nshots)
    resultS = jobS.result().get_counts(); A1B1_avg = AB_avg(resultS, nshots)
    jobE = execute(qc, backend = device, shots = nshots); job_monitor(jobE)
    A1B1_avgE = AB_avg(jobE.result().get_counts(), nshots)
    # A1B2
    qc = QuantumCircuit(2, 2)
    qc_state_prep_ = qc_state_prep(); qc.append(qc_state_prep_, [0, 1])
    qc_A1B2_ = qc_A1B2(); qc.append(qc_A1B2_, [0, 1]); qc.measure([0, 1], [0, 1])
    jobS = execute(qc, backend = simulator, shots = nshots)
    resultS = jobS.result().get_counts(); A1B2_avg = AB_avg(resultS, nshots)
    jobE = execute(qc, backend = device, shots = nshots); job_monitor(jobE)
    A1B2_avgE = AB_avg(jobE.result().get_counts(), nshots)
    # A2B1
    qc = QuantumCircuit(2, 2)
    qc_state_prep_ = qc_state_prep(); qc.append(qc_state_prep_, [0, 1])
    qc_A2B1_ = qc_A2B1(th[j]); qc.append(qc_A2B1_, [0, 1]); qc.measure([0, 1], [0, 1])
    jobS = execute(qc, backend = simulator, shots = nshots)
    resultS = jobS.result().get_counts(); A2B1_avg = AB_avg(resultS, nshots)
    jobE = execute(qc, backend = device, shots = nshots); job_monitor(jobE)
    A2B1_avgE = AB_avg(jobE.result().get_counts(), nshots)
    # A2B2
    qc = QuantumCircuit(2, 2)
    qc_state_prep_ = qc_state_prep(); qc.append(qc_state_prep_, [0, 1])
    qc_A2B2_ = qc_A2B2(); qc.append(qc_A2B2_, [0, 1]); qc.measure([0, 1], [0, 1])
    jobS = execute(qc, backend = simulator, shots = nshots)
    resultS = jobS.result().get_counts(); A2B2_avg = AB_avg(resultS, nshots)
    jobE = execute(qc, backend = device, shots = nshots); job_monitor(jobE)
    A2B2_avgE = AB_avg(jobE.result().get_counts(), nshots)
    # O
    Oavg_sim[j] = A1B1_avg + A1B2_avg + A2B1_avg - A2B2_avg
    Oavg_exp[j] = A1B1_avgE + A1B2_avgE + A2B1_avgE - A2B2_avgE

matplotlib.rcParams.update({'font.size':12})
plt.figure(figsize = (6,4), dpi = 100)
plt.plot(th, y, label = 'LHV limit')
plt.plot(th, Oavg, label = r'$\langle O\rangle_{\Psi}^{teo}$')
plt.plot(th, Oavg_sim, '*', label = r'$\langle O\rangle_{\Psi}^{sim}$')
plt.plot(th, Oavg_exp, 'o', label = r'$\langle O\rangle_{\Psi}^{exp}$')
plt.xlabel(r'$\theta$'); plt.legend(); plt.show()

def chsh(rho):
    cm = zeros(3,3); cm = corr_mat(2, 2, rho)
    evals = zeros(3,1); W = zeros(3,1); W = eVals(3, cm)
    return max(0,(sqrt(W[0]**2+W[1]**2+W[2]**2-min(W[0],W[1],W[2])**2)-1)/(sqrt(2)-1))