!pip install -U -r grading_tools/requirements.txt

from IPython.display import clear_output
clear_output()

import numpy as np; pi = np.pi
from qiskit import QuantumCircuit, Aer, execute
from qiskit.visualization import plot_histogram
from copy import deepcopy as make_copy

def prepare_hets_circuit(depth, angle1, angle2):

    hets_circ = QuantumCircuit(depth)
    hets_circ.ry(angle1, 0)
    hets_circ.rz(angle1, 0)
    hets_circ.ry(angle1, 1)
    hets_circ.rz(angle1, 1)

    for ii in range(depth):
        hets_circ.cx(0,1)
        hets_circ.ry(angle2,0)
        hets_circ.rz(angle2,0)
        hets_circ.ry(angle2,1)
        hets_circ.rz(angle2,1)
    
    return hets_circ 

hets_circuit = prepare_hets_circuit(2, pi/2, pi/2)
hets_circuit.draw()

def measure_zz_circuit(given_circuit):
    zz_meas = make_copy(given_circuit)
    zz_meas.measure_all()
    return zz_meas

zz_meas = measure_zz_circuit(hets_circuit)
zz_meas.draw()

simulator = Aer.get_backend('qasm_simulator')

result = execute(zz_meas, backend = simulator, shots=10000).result()
counts = result.get_counts(zz_meas)

plot_histogram(counts)

def measure_zz(given_circuit, num_shots = 10000):

    zz_meas = measure_zz_circuit(given_circuit)
    
    result = execute(zz_meas, backend = simulator, shots = num_shots).result()
    counts = result.get_counts(zz_meas)

    if '00' not in counts:
        counts['00'] = 0
    if '01' not in counts:
        counts['01'] = 0
    if '10' not in counts:
        counts['10'] = 0
    if '11' not in counts:
        counts['11'] = 0 

    total_counts = counts['00'] + counts['11'] + counts['01'] + counts['10']
    zz = counts['00'] + counts['11'] - counts['01'] - counts['10']
    zz = zz / total_counts
    
    return zz

zz = measure_zz(hets_circuit)
print("<ZZ> =", str(zz))

def measure_zi(given_circuit, num_shots = 10000):
    
    zz_meas = measure_zz_circuit(given_circuit)
    
    result = execute(zz_meas, backend = simulator, shots = num_shots).result()
    counts = result.get_counts(zz_meas)

    if '00' not in counts:
        counts['00'] = 0
    if '01' not in counts:
        counts['01'] = 0
    if '10' not in counts:
        counts['10'] = 0
    if '11' not in counts:
        counts['11'] = 0 

    total_counts = counts['00'] + counts['11'] + counts['01'] + counts['10']
    
    zi = counts['00'] - counts['11'] + counts['01'] - counts['10']
    zi = zi / total_counts
    
    return zi

def measure_iz(given_circuit, num_shots = 10000):
    
    zz_meas = measure_zz_circuit(given_circuit)
    
    result = execute(zz_meas, backend = simulator, shots = num_shots).result()
    counts = result.get_counts(zz_meas)

    if '00' not in counts:
        counts['00'] = 0
    if '01' not in counts:
        counts['01'] = 0
    if '10' not in counts:
        counts['10'] = 0
    if '11' not in counts:
        counts['11'] = 0 

    total_counts = counts['00'] + counts['11'] + counts['01'] + counts['10']
    
    iz = counts['00'] - counts['11'] - counts['01'] + counts['10']
    iz = iz / total_counts
    
    return iz

zi = measure_zi(hets_circuit)
print("<ZI> =", str(zi))

iz = measure_iz(hets_circuit)
print("<IZ> =", str(iz))

def measure_xx_circuit(given_circuit):
    xx_meas = make_copy(given_circuit)
    
    ### WRITE YOUR CODE BETWEEN THESE LINES - START
    
    xx_meas.h(0)
    xx_meas.h(1)
    xx_meas.measure_all()
    
    ### WRITE YOUR CODE BETWEEN THESE LINES - END

    return xx_meas

xx_meas = measure_xx_circuit(hets_circuit)
xx_meas.draw()

def measure_xx(given_circuit, num_shots = 10000):
    
    xx_meas = measure_xx_circuit(given_circuit)
    
    result = execute(xx_meas, backend = simulator, shots = num_shots).result()
    counts = result.get_counts(xx_meas)

    if '00' not in counts:
        counts['00'] = 0
    if '01' not in counts:
        counts['01'] = 0
    if '10' not in counts:
        counts['10'] = 0
    if '11' not in counts:
        counts['11'] = 0

    total_counts = counts['00'] + counts['11'] + counts['01'] + counts['10']
    xx = counts['00'] + counts['11'] - counts['01'] - counts['10']
    xx = xx / total_counts
    
    return xx

xx = measure_xx(hets_circuit)
print("<XX> =", str(xx))

def get_energy(given_circuit, num_shots = 10000):
    
    zz = measure_zz(given_circuit, num_shots = num_shots)
    iz = measure_iz(given_circuit, num_shots = num_shots)
    zi = measure_zi(given_circuit, num_shots = num_shots)
    xx = measure_xx(given_circuit, num_shots = num_shots)
    
    energy = (-1.0523732)*1 + (0.39793742)*iz + (-0.3979374)*zi + (-0.0112801)*zz + (0.18093119)*xx
    
    return energy

energy = get_energy(hets_circuit)
print("The energy of the trial state is", str(energy))

hets_circuit_plus = None
hets_circuit_minus = None

### WRITE YOUR CODE BETWEEN THESE LINES - START

hets_circuit_plus = prepare_hets_circuit(2, pi/2 + 0.1*pi/2, pi/2)
hets_circuit_minus = prepare_hets_circuit(2, pi/2 - 0.1*pi/2, pi/2)

### WRITE YOUR CODE BETWEEN THESE LINES - END

energy_plus = get_energy(hets_circuit_plus, num_shots=100000)
energy_minus = get_energy(hets_circuit_minus, num_shots=100000)

print(energy_plus, energy_minus)

name = 'Pon Rahul M'
email = 'ponrahul.21it@licet.ac.in'

### Do not change the lines below
from grading_tools import grade

grade(answer=measure_xx_circuit(hets_circuit), name=name, email=email, labid='lab9', exerciseid='ex1')
grade(answer=hets_circuit_plus, name=name, email=email, labid='lab9', exerciseid='ex2')
grade(answer=hets_circuit_minus, name=name, email=email, labid='lab9', exerciseid='ex3')

energy_plus_100, energy_plus_1000, energy_plus_10000 = 0, 0, 0
energy_minus_100, energy_minus_1000, energy_minus_10000 = 0, 0, 0

### WRITE YOUR CODE BETWEEN THESE LINES - START

energy_plus_100 = get_energy(hets_circuit_plus, num_shots = 100)
energy_minus_100 = get_energy(hets_circuit_minus, num_shots = 100)

energy_plus_1000 = get_energy(hets_circuit_plus, num_shots = 1000)
energy_minus_1000 = get_energy(hets_circuit_minus, num_shots = 1000)

energy_plus_10000 = get_energy(hets_circuit_plus, num_shots = 10000)
energy_minus_10000 = get_energy(hets_circuit_minus, num_shots = 10000)

### WRITE YOUR CODE BETWEEN THESE LINES - END

print(energy_plus_100, energy_minus_100, "difference = ", energy_minus_100 - energy_plus_100)
print(energy_plus_1000, energy_minus_1000, "difference = ", energy_minus_1000 - energy_plus_1000)
print(energy_plus_10000, energy_minus_10000, "difference = ", energy_minus_10000 - energy_plus_10000)

### WRITE YOUR CODE BETWEEN THESE LINES - START

I = np.array([
[1, 0],
[0, 1]
])

X = np.array([
[0, 1],
[1, 0]
])

Z = np.array([
[1, 0],
[0, -1]
])

h2_hamiltonian = (-1.0523732) * np.kron(I, I) + \
                 (0.39793742) * np.kron(I, Z) + \
                 (-0.3979374) * np.kron(Z, I) + \
                 (-0.0112801) * np.kron(Z, Z) + \
                 (0.18093119) * np.kron(X, X)

from numpy import linalg as LA
eigenvalues, eigenvectors = LA.eig(h2_hamiltonian)
for ii, eigenvalue in enumerate(eigenvalues):
    print(f"Eigenvector {eigenvectors[:,ii]} has energy {eigenvalue}")

exact_eigenvector = eigenvectors[:,np.argmin(eigenvalues)]
exact_eigenvalue = np.min(eigenvalues)

print()
print("Minimum energy is", exact_eigenvalue)

### WRITE YOUR CODE BETWEEN THESE LINES - END

