# install cirq
!pip install cirq==0.5 --quiet

import cirq
import numpy as np
import random

print(cirq.google.Bristlecone)

def make_quantum_teleportation_circuit(gate):
    circuit = cirq.Circuit()
    msg, alice, bob = cirq.LineQubit.range(3)

    # Creates Bell state to be shared between Alice and Bob
    circuit.append([cirq.H(alice), cirq.CNOT(alice, bob)])
    # Creates a random state for the Message
    circuit.append(gate(msg))
    # Bell measurement of the Message and Alice's entangled qubit
    circuit.append([cirq.CNOT(msg, alice), cirq.H(msg)])
    circuit.append(cirq.measure(msg, alice))
    # Uses the two classical bits from the Bell measurement to recover the
    # original quantum Message on Bob's entangled qubit
    circuit.append([cirq.CNOT(alice, bob), cirq.CZ(msg, bob)])

    return circuit


gate = cirq.SingleQubitMatrixGate(cirq.testing.random_unitary(2))
circuit = make_quantum_teleportation_circuit(gate)

print("Circuit:")
print(circuit)

sim = cirq.Simulator()

# Create qubits.
q0 = cirq.LineQubit

# Produces the message using random unitary
message = sim.simulate(cirq.Circuit.from_ops(gate(q0)))

print("Bloch Vector of Message After Random Unitary:")
# Prints the Bloch vector of the Message after the random gate
b0X, b0Y, b0Z = cirq.bloch_vector_from_state_vector(
    message.final_state, 0)
print("x: ", np.around(b0X, 4),
      "y: ", np.around(b0Y, 4),
      "z: ", np.around(b0Z, 4))

# Records the final state of the simulation
final_results = sim.simulate(circuit)

print("Bloch Sphere of Qubit 2 at Final State:")
# Prints the Bloch Sphere of Bob's entangled qubit at the final state
b2X, b2Y, b2Z = cirq.bloch_vector_from_state_vector(
    final_results.final_state, 2)
print("x: ", np.around(b2X, 4),
      "y: ", np.around(b2Y, 4),
      "z: ", np.around(b2Z, 4))

def make_oracle(q0, q1, secret_function):
    """ Gates implementing the secret function f(x)."""

    # coverage: ignore
    if secret_function[0]:
        yield [cirq.CNOT(q0, q1), cirq.X(q1)]

    if secret_function[1]:
        yield cirq.CNOT(q0, q1)


def make_deutsch_circuit(q0, q1, oracle):
    c = cirq.Circuit()

    # Initialize qubits.
    c.append([cirq.X(q1), cirq.H(q1), cirq.H(q0)])

    # Query oracle.
    c.append(oracle)

    # Measure in X basis.
    c.append([cirq.H(q0), cirq.measure(q0, key='result')])
    return c
  
# Choose qubits to use.
q0, q1 = cirq.LineQubit.range(2)

# Pick a secret 2-bit function and create a circuit to query the oracle.
secret_function = [random.randint(0,1) for _ in range(2)]
oracle = make_oracle(q0, q1, secret_function)
print('Secret function:\nf(x) = <{}>'.format(
    ', '.join(str(e) for e in secret_function)))

# Embed the oracle into a quantum circuit querying it exactly once.
circuit = make_deutsch_circuit(q0, q1, oracle)
print('Circuit:')
print(circuit)

# Simulate the circuit.
simulator = cirq.Simulator()
result = simulator.run(circuit)
print('Result of f(0)⊕f(1):')
print(result)

def make_qft(qubits):
  """Generator for the QFT on an arbitrary number of qubits. With four qubits
  the answer is
  ---H--@-------@--------@---------------------------------------------
        |       |        |
  ------@^0.5---+--------+---------H--@-------@------------------------
                |        |            |       |
  --------------@^0.25---+------------@^0.5---+---------H--@-----------
                         |                    |            |
  -----------------------@^0.125--------------@^0.25-------@^0.5---H---
  """
  # YOUR CODE HERE

def make_qft(qubits):
  """Generator for the QFT on an arbitrary number of qubits. With four qubits
  the answer is
  ---H--@-------@--------@---------------------------------------------
        |       |        |
  ------@^0.5---+--------+---------H--@-------@------------------------
                |        |            |       |
  --------------@^0.25---+------------@^0.5---+---------H--@-----------
                         |                    |            |
  -----------------------@^0.125--------------@^0.25-------@^0.5---H---
  """
  qubits = list(qubits)
  while len(qubits) > 0:
      q_head = qubits.pop(0)
      yield cirq.H(q_head)
      for i, qubit in enumerate(qubits):
          yield (cirq.CZ**(1/2**(i+1)))(qubit, q_head)

num_qubits = 4
qubits = cirq.LineQubit.range(num_qubits)

qft = cirq.Circuit.from_ops(make_qft(qubits))
print(qft)

class QFT(cirq.Gate):
  """Gate for the Quantum Fourier Transformation
  """
  
  def __init__(self, n_qubits):
    self.n_qubits = n_qubits

  def num_qubits(self):
    return self.n_qubits    
    
  def _decompose_(self, qubits):
    # YOUR CODE HERE
            
  # How should the gate look in ASCII diagrams?          
  def _circuit_diagram_info_(self, args):        
    return tuple('QFT{}'.format(i) for i in range(self.n_qubits))

class QFT(cirq.Gate):
  """Gate for the Quantum Fourier Transformation
  """
  
  def __init__(self, n_qubits):
    self.n_qubits = n_qubits

  def num_qubits(self):
    return self.n_qubits
    
  def _decompose_(self, qubits):
    """Implements the QFT on an arbitrary number of qubits. The circuit
    for num_qubits = 4 is given by
    ---H--@-------@--------@---------------------------------------------
          |       |        |
    ------@^0.5---+--------+---------H--@-------@------------------------
                  |        |            |       |
    --------------@^0.25---+------------@^0.5---+---------H--@-----------
                           |                    |            |
    -----------------------@^0.125--------------@^0.25-------@^0.5---H---
    """
    qubits = list(qubits)
    while len(qubits) > 0:
        q_head = qubits.pop(0)
        yield cirq.H(q_head)
        for i, qubit in enumerate(qubits):
            yield (cirq.CZ**(1/2**(i+1)))(qubit, q_head)
            
  # How should the gate look in ASCII diagrams?          
  def _circuit_diagram_info_(self, args):        
    return tuple('QFT{}'.format(i) for i in range(self.n_qubits))

num_qubits = 4

qubits = cirq.LineQubit.range(num_qubits)
circuit = cirq.Circuit.from_ops(QFT(num_qubits).on(*qubits))
print(circuit)

qft_test = cirq.Circuit.from_ops(make_qft(qubits))
print(qft_test)
np.testing.assert_allclose(cirq.unitary(qft_test), cirq.unitary(circuit))

class QFT_inv(cirq.Gate):
  """Gate for the inverse Quantum Fourier Transformation
  """
  
  # YOUR CODE HERE

class QFT_inv(cirq.Gate):
  """Gate for the inverse Quantum Fourier Transformation
  """
  
  def __init__(self, n_qubits):
    self.n_qubits = n_qubits

  def num_qubits(self):
    return self.n_qubits   
    
  def _decompose_(self, qubits):
    """Implements the inverse QFT on an arbitrary number of qubits. The circuit
    for num_qubits = 4 is given by
    ---H--@-------@--------@---------------------------------------------
          |       |        |
    ------@^-0.5--+--------+---------H--@-------@------------------------
                  |        |            |       |
    --------------@^-0.25--+------------@^-0.5--+---------H--@-----------
                           |                    |            |
    -----------------------@^-0.125-------------@^-0.25------@^-0.5--H---
    """
    qubits = list(qubits)
    while len(qubits) > 0:
        q_head = qubits.pop(0)
        yield cirq.H(q_head)
        for i, qubit in enumerate(qubits):
            yield (cirq.CZ**(-1/2**(i+1)))(qubit, q_head)
                      
  def _circuit_diagram_info_(self, args):        
    return tuple('QFT{}^-1'.format(i) for i in range(self.n_qubits))

num_qubits = 2

qubits = cirq.LineQubit.range(num_qubits)
circuit = cirq.Circuit.from_ops(QFT(num_qubits).on(*qubits),
                                QFT_inv(num_qubits).on(*qubits))
print(circuit)
cirq.unitary(circuit).round(2)

num_qubits = 2

qubits = cirq.LineQubit.range(num_qubits)
circuit = cirq.Circuit.from_ops(QFT(num_qubits).on(*qubits),
                                QFT_inv(num_qubits).on(*qubits[::-1])) # qubit order reversed
print(circuit)
cirq.unitary(circuit)

theta = 0.234 # Try your own
n_bits = 3 # Accuracy of the estimate for theta. Try different values.

qubits = cirq.LineQubit.range(n_bits)
u_bit = cirq.NamedQubit('u')

U = cirq.Z**(2*theta)

phase_estimator = cirq.Circuit()

phase_estimator.append(cirq.H.on_each(*qubits))
for i, bit in enumerate(qubits):
  phase_estimator.append(cirq.ControlledGate(U).on(bit,u_bit)**(2**(n_bits-1-i)))

print(phase_estimator)

phase_estimator.append(QFT_inv(n_bits).on(*qubits))
print(phase_estimator)

# Add measurements to the end of the circuit
phase_estimator.append(cirq.measure(*qubits, key='m'))

# Add gate to change initial state to |1>
phase_estimator.insert(0,cirq.X(u_bit))

print(phase_estimator)

sim = cirq.Simulator()
result = sim.run(phase_estimator, repetitions=10)

theta_estimates = np.sum(2**np.arange(n_bits)*result.measurements['m'], axis=1)/2**n_bits
print(theta_estimates)

def phase_estimation(theta, n_bits, n_reps=10):

  # YOUR CODE HERE
  
  return theta_estimates

def phase_estimation(theta, n_bits, n_reps=10):
  qubits = cirq.LineQubit.range(n_bits)
  u_bit = cirq.NamedQubit('u')

  U = cirq.Z**(2*theta) # Try out a different gate if you like

  phase_estimator = cirq.Circuit()

  phase_estimator.append(cirq.H.on_each(*qubits))
  for i, bit in enumerate(qubits):
    phase_estimator.append(cirq.ControlledGate(U).on(bit,u_bit)**(2**(n_bits-1-i)))
  phase_estimator.append(QFT_inv(n_bits).on(*qubits))
  
  # Measurement gates
  phase_estimator.append(cirq.measure(*qubits, key='m'))

  # Gates to choose initial state for the u_bit. Placing X here chooses the |1> state
  phase_estimator.insert(0,cirq.X(u_bit))

  # Code to simulate measurements
  sim = cirq.Simulator()
  result = sim.run(phase_estimator, repetitions=n_reps)
  
  # Convert measurements into estimates of theta
  theta_estimates = np.sum(2**np.arange(n_bits)*result.measurements['m'], axis=1)/2**n_bits
  
  return theta_estimates

phase_estimation(0.234, 10)

def phase_estimation(theta, n_bits, n_reps=10):
  qubits = cirq.LineQubit.range(n_bits)
  u_bit = cirq.NamedQubit('u')

  U = cirq.Z**(2*theta)

  phase_estimator = cirq.Circuit()

  phase_estimator.append(cirq.H.on_each(*qubits))
  for i, bit in enumerate(qubits):
    phase_estimator.append(cirq.ControlledGate(U).on(bit,u_bit)**(2**(n_bits-1-i))) # Could have used CZ in this example
  phase_estimator.append(QFT_inv(n_bits).on(*qubits))
  
  phase_estimator.append(cirq.measure(*qubits, key='m'))

  # Changed the X gate here to an H
  phase_estimator.insert(0,cirq.H(u_bit))

  sim = cirq.Simulator()
  result = sim.run(phase_estimator, repetitions=n_reps)
    
  theta_estimates = np.sum(2**np.arange(n_bits)*result.measurements['m'], axis=1)/2**n_bits
  
  return theta_estimates

phase_estimation(0.234,10)

def set_io_qubits(qubit_count):
    """Add the specified number of input and output qubits."""
    input_qubits = [cirq.GridQubit(i, 0) for i in range(qubit_count)]
    output_qubit = cirq.GridQubit(qubit_count, 0)
    return (input_qubits, output_qubit)


def make_oracle(input_qubits, output_qubit, x_bits):
    """Implement function {f(x) = 1 if x==x', f(x) = 0 if x!= x'}."""
    # Make oracle.
    # for (1, 1) it's just a Toffoli gate
    # otherwise negate the zero-bits.
    yield(cirq.X(q) for (q, bit) in zip(input_qubits, x_bits) if not bit)
    yield(cirq.TOFFOLI(input_qubits[0], input_qubits[1], output_qubit))
    yield(cirq.X(q) for (q, bit) in zip(input_qubits, x_bits) if not bit)


def make_grover_circuit(input_qubits, output_qubit, oracle):
    """Find the value recognized by the oracle in sqrt(N) attempts."""
    # For 2 input qubits, that means using Grover operator only once.
    c = cirq.Circuit()

    # Initialize qubits.
    c.append([
        cirq.X(output_qubit),
        cirq.H(output_qubit),
        cirq.H.on_each(*input_qubits),
    ])

    # Query oracle.
    c.append(oracle)

    # Construct Grover operator.
    c.append(cirq.H.on_each(*input_qubits))
    c.append(cirq.X.on_each(*input_qubits))
    c.append(cirq.H.on(input_qubits[1]))
    c.append(cirq.CNOT(input_qubits[0], input_qubits[1]))
    c.append(cirq.H.on(input_qubits[1]))
    c.append(cirq.X.on_each(*input_qubits))
    c.append(cirq.H.on_each(*input_qubits))

    # Measure the result.
    c.append(cirq.measure(*input_qubits, key='result'))

    return c


def bitstring(bits):
    return ''.join(str(int(b)) for b in bits)


qubit_count = 2
circuit_sample_count = 10

#Set up input and output qubits.
(input_qubits, output_qubit) = set_io_qubits(qubit_count)

#Choose the x' and make an oracle which can recognize it.
x_bits = [random.randint(0, 1) for _ in range(qubit_count)]
print('Secret bit sequence: {}'.format(x_bits))

# Make oracle (black box)
oracle = make_oracle(input_qubits, output_qubit, x_bits)

# Embed the oracle into a quantum circuit implementing Grover's algorithm.
circuit = make_grover_circuit(input_qubits, output_qubit, oracle)
print('Circuit:')
print(circuit)

# Sample from the circuit a couple times.
simulator = cirq.Simulator()
result = simulator.run(circuit, repetitions=circuit_sample_count)

frequencies = result.histogram(key='result', fold_func=bitstring)
print('Sampled results:\n{}'.format(frequencies))

# Check if we actually found the secret value.
most_common_bitstring = frequencies.most_common(1)[0][0]
print('Most common bitstring: {}'.format(most_common_bitstring))
print('Found a match: {}'.format(
    most_common_bitstring == bitstring(x_bits)))
