# install release containing NeutralAtomDevice and IonDevice classes
!pip install cirq~=0.5.0 --quiet

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

ms = cirq.Duration(nanos=10**6)
us = cirq.Duration(nanos=10**3)
neutral_device = cirq.NeutralAtomDevice(measurement_duration = 5*ms,
                                       gate_duration = 100*us,
                                       control_radius = 2,
                                       max_parallel_z = 3,
                                       max_parallel_xy = 3,
                                       max_parallel_c = 3,
                                       qubits = [cirq.GridQubit(row, col) 
                                                 for col in range(3) 
                                                 for row in range(3)])

# Single Qubit Z rotation by Pi/5 radians
neutral_device.validate_gate(cirq.Rz(np.pi/5))
# Single Qubit rotation about axis in X-Y plane Pi/3 radians from X axis by angle of Pi/7
neutral_device.validate_gate(cirq.PhasedXPowGate(phase_exponent=np.pi/3,exponent=np.pi/7))
# Controlled gate with integer exponent 
neutral_device.validate_gate(cirq.CNOT)
# Controlled Not gate with two controls
neutral_device.validate_gate(cirq.TOFFOLI)

#Controlled gate with non-integer exponent (rotation angle must be a multiple of pi)
neutral_device.validate_gate(cirq.CZ**1.5)

# Hadamard gates rotate about the X-Z axis, which isn't compatable with our single qubit rotations
neutral_device.validate_gate(cirq.H)

ms = cirq.Duration(nanos=10**6)
us = cirq.Duration(nanos=10**3)
neutral_device = cirq.NeutralAtomDevice(measurement_duration = 5*ms,
                                       gate_duration = 100*us,
                                       control_radius = 2,
                                       max_parallel_z = 3,
                                       max_parallel_xy = 3,
                                       max_parallel_c = 3,
                                       qubits = [cirq.GridQubit(row, col) 
                                                 for col in range(3) 
                                                 for row in range(3)])

# Moment/Circuit Examples
moment_circ = cirq.Circuit(device=neutral_device)
qubits = [cirq.GridQubit(row, col) for col in range(3) for row in range(3)]

# Three qubits affected by a Z gate in parallel with Three qubits affected
# by an XY gate
operation_list_one = cirq.Z.on_each(*qubits[:3])+cirq.X.on_each(*qubits[3:6])
valid_moment_one = cirq.Moment(operation_list_one)
moment_circ.append(valid_moment_one)

# A TOFFOLI gate on three qubits that are close enough to eachother
operation_list_two = [cirq.TOFFOLI.on(*qubits[:3])]
valid_moment_two = cirq.Moment(operation_list_two)
moment_circ.append(valid_moment_two)


print(moment_circ)

global_circuit = cirq.Circuit(device=neutral_device)
global_list_of_operations = cirq.X.on_each(*qubits)
global_circuit.append(global_list_of_operations)
print(global_circuit)

global_moment_circuit = cirq.Circuit(device=neutral_device)
global_moment = cirq.Moment(cirq.X.on_each(*qubits))
global_moment_circuit.append(global_moment)
print(global_moment_circuit)

parallel_gate_op_circuit = cirq.Circuit(device=neutral_device)
parallel_gate_op = cirq.ParallelGateOperation(cirq.X,qubits)
parallel_gate_op_circuit.append(parallel_gate_op)
print(parallel_gate_op_circuit)

def oracle(qubits, key_bits):
    yield (cirq.X(q) for (q, bit) in zip(qubits, key_bits) if not bit)
    yield cirq.CCZ(*qubits)
    yield (cirq.X(q) for (q, bit) in zip(qubits, key_bits) if not bit)

# Try changing the key to see the relationship between
# the placement of the X gates and the key
key = (1, 0, 1)
qubits = [cirq.GridQubit(0,col) for col in range(3)]
oracle_example_circuit = cirq.Circuit().from_ops(oracle(qubits,key))
print(oracle_example_circuit)

def diffusion_operator(qubits):
    yield cirq.H.on_each(*qubits)
    yield cirq.X.on_each(*qubits)
    yield cirq.CCZ(*qubits)
    yield cirq.X.on_each(*qubits)
    yield cirq.H.on_each(*qubits)
    
qubits = [cirq.GridQubit(0,col) for col in range(3)]
diffusion_circuit = cirq.Circuit().from_ops(diffusion_operator(qubits))
print(diffusion_circuit)

def initial_hadamards(qubits):
    yield cirq.H.on_each(*qubits)

uncompiled_circuit = cirq.Circuit()
key = (1,0,1)
qubits = [cirq.GridQubit(0,0),cirq.GridQubit(0,1),cirq.GridQubit(0,2)]
uncompiled_circuit.append(initial_hadamards(qubits))
uncompiled_circuit.append(oracle(qubits,key))
uncompiled_circuit.append(diffusion_operator(qubits))
uncompiled_circuit.append(oracle(qubits,key))
uncompiled_circuit.append(diffusion_operator(qubits))
print(uncompiled_circuit)

def neutral_atom_initial_step(qubits):
    yield cirq.ParallelGateOperation(cirq.Y**(1/2), qubits)
    
def neutral_atom_diffusion_operator(qubits):
    yield cirq.ParallelGateOperation(cirq.Y**(1/2), qubits)
    yield cirq.CCZ(*qubits)
    yield cirq.ParallelGateOperation(cirq.Y**(-1/2), qubits)

    
ms = cirq.Duration(nanos=10**6)
us = cirq.Duration(nanos=10**3)
qubits = [cirq.GridQubit(row, col) for col in range(3) for row in range(1)]
three_qubit_device = cirq.NeutralAtomDevice(measurement_duration = 5*ms,
                                           gate_duration = us,
                                           control_radius = 2,
                                           max_parallel_z = 3,
                                           max_parallel_xy = 3,
                                           max_parallel_c = 3,
                                           qubits=qubits)
key = (0,1,0)
compiled_grover_circuit = cirq.Circuit(device=three_qubit_device)
compiled_grover_circuit.append(neutral_atom_initial_step(qubits))
compiled_grover_circuit.append(oracle(qubits,key))
compiled_grover_circuit.append(neutral_atom_diffusion_operator(qubits))
compiled_grover_circuit.append(oracle(qubits,key))
compiled_grover_circuit.append(neutral_atom_diffusion_operator(qubits))
print(compiled_grover_circuit)

def grover_circuit_with_n_repetitions(n, key):
    ms = cirq.Duration(nanos=10**6)
    us = cirq.Duration(nanos=10**3)
    qubits = [cirq.GridQubit(row, col) for col in range(3) for row in range(1)]
    three_qubit_device = cirq.NeutralAtomDevice(measurement_duration = 5*ms,
                                               gate_duration = us,
                                               control_radius = 2,
                                               max_parallel_z = 3,
                                               max_parallel_xy = 3,
                                               max_parallel_c = 3,
                                               qubits=qubits)
    grover_circuit = cirq.Circuit(device=three_qubit_device)
    grover_circuit.append(neutral_atom_initial_step(qubits))
    for repetition in range(n):
      grover_circuit.append(oracle(qubits,key))
      grover_circuit.append(neutral_atom_diffusion_operator(qubits))
      
    return grover_circuit
      
success_probabilities = []
key = (0,1,1)
N = 2**3
#Convert key from binary to a base 10 number
diag = sum(2**(2-count) for (count, val) in enumerate(key) if val)
num_points = 10
for repetitions in range(num_points):
    test_circuit = grover_circuit_with_n_repetitions(repetitions, key)
    sim = cirq.Simulator()
    result = sim.simulate(test_circuit)
    rho = result.density_matrix_of(qubits)
    success_probabilities.append(np.real(rho[diag][diag]))

plt.scatter(range(num_points), success_probabilities, label="Simulation")
x = np.linspace(0, num_points, 1000)
y = np.sin((2*x+1)*np.arcsin(1/np.sqrt(N)))**2
plt.plot(x, y, label="Theoretical Curve")
plt.title("Probability of Success Vs. Number of Oracle-Diffusion Operators")
plt.ylabel("Probability of Success")
plt.xlabel("Number of Times Oracle and Diffusion Operators are Applied")
plt.legend(loc='upper right')
plt.show()