

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

import cirq
import numpy as np
import cirq.ion as ci
from cirq import Simulator
import itertools
import random
### number of qubits 
qubit_num = 5
### define your qubits as line qubits for a linear ion trap
qubit_list = cirq.LineQubit.range(qubit_num)
### make your ion trap device with desired gate times and qubits
us = 1000*cirq.Duration(nanos=1)
ion_device = ci.IonDevice(measurement_duration=100*us,
                        twoq_gates_duration=200*us,
                        oneq_gates_duration=10*us,
                        qubits=qubit_list)

# Single Qubit Z rotation by Pi/5 radians
ion_device.validate_gate(cirq.Rz(np.pi/5))
# Single Qubit X rotation by Pi/7 radians 
ion_device.validate_gate(cirq.Rx(np.pi/7))
# Molmer-Sorensen gate by Pi/4 
ion_device.validate_gate(cirq.MS(np.pi/4))

#Controlled gate with non-integer exponent (rotation angle must be a multiple of pi)
ion_device.validate_gate(cirq.TOFFOLI)

### length of your hidden string
string_length = qubit_num-1
### generate all possible strings of length string_length, and randomly choose one as your hidden string
all_strings = ["".join(seq) for seq in itertools.product("01", repeat=string_length)]
hidden_string = random.choice(all_strings)


### make the circuit for BV with clifford gates
circuit = cirq.Circuit()
circuit.append([cirq.X(qubit_list[qubit_num-1])])
for i in range(qubit_num):
    circuit.append([cirq.H(qubit_list[i])])
for i in range(qubit_num-1):
    if hidden_string[i] == '1':
        circuit.append([cirq.CNOT(qubit_list[i], qubit_list[qubit_num-1])])
for i in range(qubit_num - 1):
    circuit.append([cirq.H(qubit_list[i])])
    circuit.append([cirq.measure(qubit_list[i])])

print("Doing Bernstein-Vazirani algorithm with hidden string",
      hidden_string, "\n")
print("Clifford Circuit:\n")
print(circuit, "\n")

### convert the clifford circuit into circuit with ion trap native gates
ion_circuit = ion_device.decompose_circuit(circuit)
print(repr(ion_circuit))
print("Iontrap Circuit: \n", ion_circuit, "\n")

### convert the clifford circuit into circuit with ion trap native gates
ion_circuit = ion_device.decompose_circuit(circuit)
optimized_ion_circuit=cirq.merge_single_qubit_gates_into_phased_x_z(ion_circuit)
print("Iontrap Circuit: \n", ion_circuit, "\n")

### run the ion trap circuit
simulator = Simulator()
clifford_result = simulator.run(circuit)
result = simulator.run(ion_circuit)

measurement_results = ''
for i in range(qubit_num-1):
    if result.measurements[str(i)][0][0]:
        measurement_results += '1'
    else:
        measurement_results += '0'

print("Hidden string is:", hidden_string)
print("Measurement results are:", measurement_results)
print("Found answer using Bernstein-Vazirani:", hidden_string == measurement_results)


