from hide_toggle import hide_toggle
#usage: 
#1 create a cell with: hide_toggle(for_next=True)
#2 put the commented solution in the next cell

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

# definition of 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]])

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()

#solution

hide_toggle(for_next=True)

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 = 8

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()

#solution

hide_toggle(for_next=True)

# commutation function
def comm_check(a,b):
    aa = np.kron(a,a)
    bb = np.kron(b,b)
    return (np.dot(aa,bb) - np.dot(bb,aa))

comm_check(sx,sy)

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')

#solution

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))

#Solution: try to write the circuit and the quantum evolution,
#hints: start with:
#for k in range(len(q_times)):
    #delta = 
    #define QuantumRegister, ClassicalRegister and QuantumCircuit (Heis2)
    #define initial state
    #define each part of the circuit: ZZ, YY, XX
    #measurement
    
    # 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
    # Run the quantum algorithm and colect the result, counts
    

hide_toggle(for_next=True)

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.BasicAer.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: all these cells 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]) )


# SOLUTION:
#sxsx12 = np.kron(sx, np.kron(sx,idt))
#sysy12 = 
#szsz12 = 

#sxsx23 = 
#sysy23 = 
#szsz23 = 

#H12 = J12 * ( sxsx12 + sysy12 + szsz12 )
#H23 = 
#H =  H12 + H23

hide_toggle(for_next=True)

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))

#SOLUTION:
#for i in range(len(c_times)):
    #t = c_times[i]
    #psi_t = 
    #sz1_t[i] = 
    #sz2_t[i] =
    #sz3_t[i] =

hide_toggle(for_next=True)

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
hide_toggle(for_next=True)

### 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.BasicAer.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()

Heis3.count_ops()

import Qconfig
from qiskit.tools.monitor import job_monitor

qk.IBMQ.enable_account(Qconfig.APItoken, **Qconfig.config)

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 = qk.IBMQ.get_backend('ibmqx4')
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