import qiskit as qk
import numpy as np
from scipy.linalg import expm
import matplotlib.pyplot as plt
import math

# Single qubit operators

sx = np.array([[0.0, 1.0],[1.0, 0.0]])
sy = np.array([[0.0, -1.0*1j],[1.0*1j, 0.0]])
sz = np.array([[1.0, 0.0],[0.0, -1.0]])
idt = np.array([[1.0, 0.0],[0.0, 1.0]])

import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)

psi0 = np.array([1.0, 0.0])
thetas = np.linspace(0,4*math.pi,200)
avg_sx_tot = np.zeros(len(thetas))

for i in range(len(thetas)):
    psi_theta = expm(-1j*0.5*thetas[i]*(sx+sz)/math.sqrt(2)).dot(psi0)
    avg_sx_tot[i] = np.real(psi_theta.conjugate().transpose().dot(sx.dot(psi_theta)))

plt.plot(thetas,avg_sx_tot)
plt.xlabel(r'$\theta$')
plt.ylabel(r'$\langle\sigma_x\rangle_\theta$')
plt.show()

avg_sx_zx = np.zeros(len(thetas))
avg_sx_xz = np.zeros(len(thetas))

for i in range(len(thetas)):
    psi_theta_zx = expm(-1j*0.5*thetas[i]*(sz)/math.sqrt(2)).dot(expm(-1j*0.5*thetas[i]*(sx)/math.sqrt(2)).dot(psi0))
    psi_theta_xz = expm(-1j*0.5*thetas[i]*(sx)/math.sqrt(2)).dot(expm(-1j*0.5*thetas[i]*(sz)/math.sqrt(2)).dot(psi0))
    avg_sx_zx[i] = np.real(psi_theta_zx.conjugate().transpose().dot(sx.dot(psi_theta_zx)))
    avg_sx_xz[i] = np.real(psi_theta_xz.conjugate().transpose().dot(sx.dot(psi_theta_xz)))

plt.plot(thetas,avg_sx_tot)
plt.plot(thetas,avg_sx_zx)
plt.plot(thetas,avg_sx_xz)
plt.xlabel(r'$\theta$')
plt.ylabel(r'$\langle\sigma_x\rangle_\theta$')
plt.legend(['Around x = z', 'x first', 'z first'],loc=1)
plt.show()

# Try this with e.g. ntrot = 1, 5, 10, 50. 
# You can also try to do sx and sz slices in the reverse order: both choices will become good approximations 
# for large n

ntrot = 10

avg_sx_n = np.zeros(len(thetas))
for i in range(len(thetas)):
    rot = expm(-1j*0.5*thetas[i]*(sx)/(ntrot*math.sqrt(2))).dot(expm(-1j*0.5*thetas[i]*(sz)/(ntrot*math.sqrt(2))))
    for j in range(ntrot-1):
        rot = expm(-1j*0.5*thetas[i]*(sx)/(ntrot*math.sqrt(2))).dot(expm(-1j*0.5*thetas[i]*(sz)/(ntrot*math.sqrt(2)))).dot(rot)
    psi_theta_n = rot.dot(psi0)
    avg_sx_n[i] = np.real(psi_theta_n.conjugate().transpose().dot(sx.dot(psi_theta_n)))
    
plt.plot(thetas,avg_sx_tot)
plt.plot(thetas,avg_sx_n,'--')
plt.xlabel(r'$\theta$')
plt.ylabel(r'$\langle\sigma_x\rangle_\theta$')
plt.legend(['Exact', 'ntrot = ' + str(ntrot)],loc=1)
plt.show()

delta = 0.1
qr = qk.QuantumRegister(2,name='qr')
zz_example = qk.QuantumCircuit(qr)

zz_example.cx(qr[0],qr[1])
zz_example.u1(2*delta,qr[1])
zz_example.cx(qr[0],qr[1])

zz_example.draw(output='mpl')

J = 1
c_times = np.linspace(0,0.5*math.pi/abs(J),1000)
q_times = np.linspace(0,0.5*math.pi/abs(J),10)

### Classical simulation of the Heisenberg dimer model

psi0 = np.kron( np.array([0,1]), np.array([1,0]) )
H = J * ( np.kron(sx,sx) + np.kron(sy,sy) + np.kron(sz,sz) )

sz1_t = np.zeros(len(c_times))
sz2_t = np.zeros(len(c_times))

sz1 = np.kron(sz,idt)
sz2 = np.kron(idt,sz)

for i in range(len(c_times)):
    
    t = c_times[i]
    psi_t = expm(-1j*H*t).dot(psi0)
    sz1_t[i] = np.real(psi_t.conjugate().transpose().dot(sz1.dot(psi_t)))
    sz2_t[i] = np.real(psi_t.conjugate().transpose().dot(sz2.dot(psi_t)))

### Digital quantum simulation of the Heisenberg dimer model using qiskit

nshots = 1024

sz1q_t = np.zeros(len(q_times))
sz2q_t = np.zeros(len(q_times))

for k in range(len(q_times)):
    
    delta = J*q_times[k]
    
    qr = qk.QuantumRegister(2,name='qr')
    cr = qk.ClassicalRegister(2,name='cr')
    Heis2 = qk.QuantumCircuit(qr,cr)
    
    # Initial state preparation
    
    Heis2.x(qr[0])
    
    # ZZ
    
    Heis2.cx(qr[0],qr[1])
    Heis2.u1(-2*delta,qr[1])
    Heis2.cx(qr[0],qr[1])
    
    # YY
    
    Heis2.u3(math.pi/2, -math.pi/2, math.pi/2, qr[0])
    Heis2.u3(math.pi/2, -math.pi/2, math.pi/2, qr[1])
    Heis2.cx(qr[0],qr[1])
    Heis2.u1(-2*delta,qr[1])
    Heis2.cx(qr[0],qr[1])
    Heis2.u3(-math.pi/2, -math.pi/2, math.pi/2, qr[0])
    Heis2.u3(-math.pi/2, -math.pi/2, math.pi/2, qr[1])
    
    # XX
    
    Heis2.u3(-math.pi/2, 0.0, 0.0, qr[0])
    Heis2.u3(-math.pi/2, 0.0, 0.0, qr[1])
    Heis2.cx(qr[0],qr[1])
    Heis2.u1(-2*delta,qr[1])
    Heis2.cx(qr[0],qr[1])
    Heis2.u3(math.pi/2, 0.0, 0.0, qr[0])
    Heis2.u3(math.pi/2, 0.0, 0.0, qr[1])
    
    # measure
    
    Heis2.measure(qr,cr)
    
    # Run the quantum algorithm
    
    backend = qk.Aer.get_backend('qasm_simulator')
    job = qk.execute(Heis2, backend, shots=nshots)
    result = job.result()
    counts = result.get_counts()

    # Post processing of outcomes to get sz expectation values
    
    sz1q = 0
    sz2q = 0
    for key,value in counts.items():
        if key == '00':
            sz1q += value
            sz2q += value
        elif key == '01':
            sz1q -= value
            sz2q += value
        elif key == '10':
            sz1q += value
            sz2q -= value
        elif key == '11':
            sz1q -= value
            sz2q -= value
    
    sz1q_t[k] = sz1q/nshots
    sz2q_t[k] = sz2q/nshots


plt.plot(abs(J)*c_times,0.5*sz1_t,'b--')
plt.plot(abs(J)*c_times,0.5*sz2_t,'c')
plt.plot(abs(J)*q_times,0.5*sz1q_t,'rd')
plt.plot(abs(J)*q_times,0.5*sz2q_t,'ko')
plt.legend(['sz1','sz2','sz1q','sz2q'])
plt.xlabel(r'$\delta = |J|t$')
plt.show()


# WARNING: this cell can take a few minutes to run!

ntrotter = 5

J12 = 1
J23 = 1
c_times = np.linspace(0,math.pi/abs(J12),1000)
q_times = np.linspace(0,math.pi/abs(J12),20)

### Classical simulation of the Heisenberg trimer model

psi0 = np.kron( np.kron( np.array([0,1]), np.array([1,0]) ) , np.array([1,0]) )

sxsx12 = np.kron(sx, np.kron(sx,idt))
sysy12 = np.kron(sy, np.kron(sy,idt))
szsz12 = np.kron(sz, np.kron(sz,idt))

sxsx23 = np.kron(idt, np.kron(sx,sx))
sysy23 = np.kron(idt, np.kron(sy,sy))
szsz23 = np.kron(idt, np.kron(sz,sz))

H12 = J12 * ( sxsx12 + sysy12 + szsz12 )
H23 = J23 * ( sxsx23 + sysy23 + szsz23 )
H =  H12 + H23

sz1_t = np.zeros(len(c_times))
sz2_t = np.zeros(len(c_times))
sz3_t = np.zeros(len(c_times))

sz1 = np.kron(sz, np.kron(idt,idt))
sz2 = np.kron(idt, np.kron(sz,idt))
sz3 = np.kron(idt, np.kron(idt,sz))

for i in range(len(c_times)):
    
    t = c_times[i]
    psi_t = expm(-1j*H*t).dot(psi0)
    sz1_t[i] = np.real(psi_t.conjugate().transpose().dot(sz1.dot(psi_t)))
    sz2_t[i] = np.real(psi_t.conjugate().transpose().dot(sz2.dot(psi_t)))
    sz3_t[i] = np.real(psi_t.conjugate().transpose().dot(sz3.dot(psi_t)))
    
### Classical simulation of the Heisenberg trimer model WITH SUZUKI TROTTER DIGITALIZATION

sz1st_t = np.zeros(len(c_times))
sz2st_t = np.zeros(len(c_times))
sz3st_t = np.zeros(len(c_times))

for i in range(len(c_times)):
    
    t = c_times[i]
    Ust = expm(-1j*H23*t/ntrotter).dot(expm(-1j*H12*t/ntrotter))
    for j in range(ntrotter-1):
        Ust = expm(-1j*H23*t/ntrotter).dot(expm(-1j*H12*t/ntrotter)).dot(Ust)
    psi_t = Ust.dot(psi0)
    sz1st_t[i] = np.real(psi_t.conjugate().transpose().dot(sz1.dot(psi_t)))
    sz2st_t[i] = np.real(psi_t.conjugate().transpose().dot(sz2.dot(psi_t)))
    sz3st_t[i] = np.real(psi_t.conjugate().transpose().dot(sz3.dot(psi_t)))

    
### Digital quantum simulation of the Heisenberg model using qiskit

nshots = 1024

sz1q_t = np.zeros(len(q_times))
sz2q_t = np.zeros(len(q_times))
sz3q_t = np.zeros(len(q_times))

for k in range(len(q_times)):
    
    delta12n = J12*q_times[k]/ntrotter
    delta23n = J23*q_times[k]/ntrotter
    
    qr = qk.QuantumRegister(3,name='qr')
    cr = qk.ClassicalRegister(3,name='cr')
    Heis3 = qk.QuantumCircuit(qr,cr)
    
    # Initial state preparation
    
    Heis3.x(qr[0])
    
    for n in range(ntrotter):
        
        # 1-2 bond mapped on qubits 0 and 1 in the quantum register
        
        # ZZ
    
        Heis3.cx(qr[0],qr[1])
        Heis3.u1(-2*delta12n,qr[1])
        Heis3.cx(qr[0],qr[1])
    
        # YY
    
        Heis3.u3(math.pi/2, -math.pi/2, math.pi/2, qr[0])
        Heis3.u3(math.pi/2, -math.pi/2, math.pi/2, qr[1])
        Heis3.cx(qr[0],qr[1])
        Heis3.u1(-2*delta12n,qr[1])
        Heis3.cx(qr[0],qr[1])
        Heis3.u3(-math.pi/2, -math.pi/2, math.pi/2, qr[0])
        Heis3.u3(-math.pi/2, -math.pi/2, math.pi/2, qr[1])
    
        # XX
    
        Heis3.u3(-math.pi/2, 0.0, 0.0, qr[0])
        Heis3.u3(-math.pi/2, 0.0, 0.0, qr[1])
        Heis3.cx(qr[0],qr[1])
        Heis3.u1(-2*delta12n,qr[1])
        Heis3.cx(qr[0],qr[1])
        Heis3.u3(math.pi/2, 0.0, 0.0, qr[0])
        Heis3.u3(math.pi/2, 0.0, 0.0, qr[1])
        
        # 2-3 bond mapped on qubits 1 and 2 in the quantum register
        
        # ZZ
    
        Heis3.cx(qr[1],qr[2])
        Heis3.u1(-2*delta12n,qr[2])
        Heis3.cx(qr[1],qr[2])
    
        # YY
    
        Heis3.u3(math.pi/2, -math.pi/2, math.pi/2, qr[1])
        Heis3.u3(math.pi/2, -math.pi/2, math.pi/2, qr[2])
        Heis3.cx(qr[1],qr[2])
        Heis3.u1(-2*delta12n,qr[2])
        Heis3.cx(qr[1],qr[2])
        Heis3.u3(-math.pi/2, -math.pi/2, math.pi/2, qr[1])
        Heis3.u3(-math.pi/2, -math.pi/2, math.pi/2, qr[2])
    
        # XX
    
        Heis3.u3(-math.pi/2, 0.0, 0.0, qr[1])
        Heis3.u3(-math.pi/2, 0.0, 0.0, qr[2])
        Heis3.cx(qr[1],qr[2])
        Heis3.u1(-2*delta12n,qr[2])
        Heis3.cx(qr[1],qr[2])
        Heis3.u3(math.pi/2, 0.0, 0.0, qr[1])
        Heis3.u3(math.pi/2, 0.0, 0.0, qr[2])
    
    # measure
    
    Heis3.measure(qr,cr)
    
    # Run the quantum algorithm
    
    backend = qk.Aer.get_backend('qasm_simulator')
    job = qk.execute(Heis3, backend, shots=nshots)
    result = job.result()
    counts = result.get_counts()

    # Post processing of outcomes to get sz expectation values

    sz1q = 0
    sz2q = 0
    sz3q = 0
    for key,value in counts.items():
        if key == '000':
            sz1q += value
            sz2q += value
            sz3q += value
        elif key == '001':
            sz1q -= value
            sz2q += value
            sz3q += value
        elif key == '010':
            sz1q += value
            sz2q -= value
            sz3q += value
        elif key == '011':
            sz1q -= value
            sz2q -= value
            sz3q += value
        elif key == '100':
            sz1q += value
            sz2q += value
            sz3q -= value
        elif key == '101':
            sz1q -= value
            sz2q += value
            sz3q -= value
        elif key == '110':
            sz1q += value
            sz2q -= value
            sz3q -= value
        elif key == '111':
            sz1q -= value
            sz2q -= value
            sz3q -= value
    
    sz1q_t[k] = sz1q/nshots
    sz2q_t[k] = sz2q/nshots
    sz3q_t[k] = sz3q/nshots

fig = plt.figure()
ax = plt.subplot(111)
ax.plot(abs(J12)*c_times,0.5*sz1_t,'b', label='sz1 full')
ax.plot(abs(J12)*c_times,0.5*sz2_t,'r', label='sz2 full')
ax.plot(abs(J12)*c_times,0.5*sz3_t,'g', label='sz3 full')
ax.plot(abs(J12)*c_times,0.5*sz1st_t,'b--',label='sz1 n = ' + str(ntrotter) + ' (classical)')
ax.plot(abs(J12)*c_times,0.5*sz2st_t,'r--', label='sz2 n = ' + str(ntrotter) + ' (classical)')
ax.plot(abs(J12)*c_times,0.5*sz3st_t,'g--', label='sz3 n = ' + str(ntrotter) + ' (classical)')
ax.plot(abs(J12)*q_times,0.5*sz1q_t,'b*',label='sz1 n = ' + str(ntrotter) + ' (quantum)')
ax.plot(abs(J12)*q_times,0.5*sz2q_t,'ro', label='sz2 n = ' + str(ntrotter) + ' (quantum)')
ax.plot(abs(J12)*q_times,0.5*sz3q_t,'gd', label='sz3 n = ' + str(ntrotter) + ' (quantum)')
chartBox = ax.get_position()
ax.set_position([chartBox.x0, chartBox.y0, chartBox.width*1, chartBox.height])
ax.legend(loc='upper center', bbox_to_anchor=(1.3, 0.8), shadow=True, ncol=1)
plt.xlabel(r'$\delta_{12} = |J_{12}|t$')
plt.show()

qr = qk.QuantumRegister(1,name='qr')
cr = qk.ClassicalRegister(1,name='cr')
x_meas = qk.QuantumCircuit(qr,cr)

# Here you can add any quantum circuit generating the quantum state
# that you with to measure at the end e.g.
#
# x_meas.u3(a,b,c,qr[0])
#


# Measurement in x direction
x_meas.h(qr[0])
x_meas.measure(qr,cr)

x_meas.draw(output='mpl')

nshots = 8192
backend = qk.Aer.get_backend('qasm_simulator')
job = qk.execute(x_meas, backend, shots=nshots)
result = job.result()
counts = result.get_counts()
print(counts)

average_sigma_x = 0
for key,value in counts.items():
        if key == '0':
            average_sigma_x += value
        elif key == '1':
            average_sigma_x -= value

average_sigma_x = average_sigma_x/nshots
print('Average of sx = ', average_sigma_x)

theta_a = np.linspace(0.0, 2*math.pi, 21)     # Some angles for the a vector, with respect to the z axis
theta_b = 0.0                                 # If the b vector is the z axis -> angle = 0
theta_c = math.pi/2                           # If the vector c is the x axis -> angle = pi/2 with respect to z

# Quantum version

nshots = 1024

Pab = np.zeros(len(theta_a))
Pac = np.zeros(len(theta_a))

# Quantum circuit to measure Pbc

qr = qk.QuantumRegister(2, name='qr')
cr = qk.ClassicalRegister(2, name='cr')
qc = qk.QuantumCircuit(qr, cr)
    
## Initialize Bell state
qc.x(qr[0])
qc.x(qr[1])
qc.h(qr[0])
qc.cx(qr[0],qr[1])
    
## Pre measurement unitaries to select the basis along a (for the first qubit) and b (for the second qubit)
qc.u3(theta_b,0,0,qr[0])
qc.u3(theta_c,0,0,qr[1])

## Measures
qc.measure(qr[0], cr[0])
qc.measure(qr[1], cr[1])
    
## Running and saving result
backend = qk.Aer.get_backend('qasm_simulator')
job = qk.execute(qc, backend, shots=nshots)
result = job.result()
counts = result.get_counts()
    
out = 0
for key,value in counts.items():
    if (key == '00') or (key == '11'):
        out += value
    elif (key == '01') or (key == '10'):
        out -= value
    
Pbc = out/nshots

# Now we compute Pab and Pac for different a vectors

for i in range(len(theta_a)):
    
    ta = theta_a[i]
    
    # Quantum circuit to measure Pab with the given a

    qr = qk.QuantumRegister(2, name='qr')
    cr = qk.ClassicalRegister(2, name='cr')
    qc = qk.QuantumCircuit(qr, cr)
    
    ## Initialize Bell state
    qc.x(qr[0])
    qc.x(qr[1])
    qc.h(qr[0])
    qc.cx(qr[0],qr[1])
    
    ## Pre measurement unitaries to select the basis along a (for the first qubit) and b (for the second qubit)
    qc.u3(ta,0,0,qr[0])
    qc.u3(theta_b,0,0,qr[1])

    ## Measures
    qc.measure(qr[0], cr[0])
    qc.measure(qr[1], cr[1])
    
    ## Running and saving result
    backend = qk.Aer.get_backend('qasm_simulator')
    job = qk.execute(qc, backend, shots=nshots)
    result = job.result()
    counts = result.get_counts()
    
    out = 0
    for key,value in counts.items():
        if (key == '00') or (key == '11'):
            out += value
        elif (key == '01') or (key == '10'):
            out -= value
    
    Pab[i] = out/nshots
    
    # Quantum circuit to measure Pac with the given a

    qr = qk.QuantumRegister(2, name='qr')
    cr = qk.ClassicalRegister(2, name='cr')
    qc = qk.QuantumCircuit(qr, cr)
    
    ## Initialize Bell state
    qc.x(qr[0])
    qc.x(qr[1])
    qc.h(qr[0])
    qc.cx(qr[0],qr[1])
    
    ## Pre measurement unitaries to select the basis along a (for the first qubit) and c (for the second qubit)
    qc.u3(ta,0,0,qr[0])
    qc.u3(theta_c,0,0,qr[1])

    ## Measures
    qc.measure(qr[0], cr[0])
    qc.measure(qr[1], cr[1])
    
    ## Running and saving result
    backend = qk.Aer.get_backend('qasm_simulator')
    job = qk.execute(qc, backend, shots=nshots)
    result = job.result()
    counts = result.get_counts()
    
    out = 0
    for key,value in counts.items():
        if (key == '00') or (key == '11'):
            out += value
        elif (key == '01') or (key == '10'):
            out -= value
    
    Pac[i] = out/nshots
    
quantum_outs = Pab-Pac

# Classical calculations

thetas = np.linspace(0,2*math.pi,100)
classical_outs = np.zeros(len(thetas))

b_vec = np.array([math.sin(theta_b),0.0,math.cos(theta_b)])
c_vec = np.array([math.sin(theta_c),0.0,math.cos(theta_c)])

Bell_limit = (1 - b_vec.dot(c_vec))*np.ones(len(thetas))

for k in range(len(thetas)):
    a_vec = np.array([math.sin(thetas[k]),0.0,math.cos(thetas[k])])
    classical_outs[k] = -a_vec.dot(b_vec) + a_vec.dot(c_vec)
    
fig = plt.figure()
ax = plt.subplot(111)
ax.plot(thetas,classical_outs,'b', label='classical computation')
ax.plot(theta_a,quantum_outs,'r*', label='quantum algorithm')
ax.plot(thetas,Bell_limit,'--k', label='Bell bound')           # If the data go beyond these boundaries,
ax.plot(thetas,-Bell_limit,'--k', label='Bell bound')          # the inequality is violated
ax.plot(thetas,(1+Pbc)*np.ones(len(thetas)),'--c', label='Bell bound (quantum)')
ax.plot(thetas,-(1+Pbc)*np.ones(len(thetas)),'--c', label='Bell bound (quantum)')
chartBox = ax.get_position()
ax.set_position([chartBox.x0, chartBox.y0, chartBox.width*1, chartBox.height])
ax.legend(loc='upper center', bbox_to_anchor=(1.3, 0.8), shadow=True, ncol=1)
plt.xlabel(r'$\theta_a$')
plt.show()

from qiskit.tools.monitor import job_monitor

my_token = ''
ibmq_provider = qk.IBMQ.enable_account(my_token)

delta = 0.5*math.pi
    
qr = qk.QuantumRegister(2,name='qr')
cr = qk.ClassicalRegister(2,name='cr')
Heis2 = qk.QuantumCircuit(qr,cr)
    
# Initial state preparation
    
Heis2.x(qr[0])
    
# ZZ
    
Heis2.cx(qr[0],qr[1])
Heis2.u1(-2*delta,qr[1])
Heis2.cx(qr[0],qr[1])
    
# YY
    
Heis2.u3(math.pi/2, -math.pi/2, math.pi/2, qr[0])
Heis2.u3(math.pi/2, -math.pi/2, math.pi/2, qr[1])
Heis2.cx(qr[0],qr[1])
Heis2.u1(-2*delta,qr[1])
Heis2.cx(qr[0],qr[1])
Heis2.u3(-math.pi/2, -math.pi/2, math.pi/2, qr[0])
Heis2.u3(-math.pi/2, -math.pi/2, math.pi/2, qr[1])
    
# XX
    
Heis2.u3(-math.pi/2, 0.0, 0.0, qr[0])
Heis2.u3(-math.pi/2, 0.0, 0.0, qr[1])
Heis2.cx(qr[0],qr[1])
Heis2.u1(-2*delta,qr[1])
Heis2.cx(qr[0],qr[1])
Heis2.u3(math.pi/2, 0.0, 0.0, qr[0])
Heis2.u3(math.pi/2, 0.0, 0.0, qr[1])
    
# measure
    
Heis2.measure(qr,cr)

my_backend = ibmq_provider.get_backend('ibmqx2')
job = qk.execute(Heis2, backend=my_backend, shots=1024) 

job_monitor(job, interval=5)

result = job.result()
counts = result.get_counts()
print(counts)
plot_histogram(counts)

my_backend.configuration().coupling_map # In the output, the first entry in a pair [a,b] is a control, second is the 
                                        # corresponding target for a CNOT