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

import cirq
import numpy as np
import sympy
import matplotlib.pyplot as plt

print(cirq.google.Bristlecone)

a = cirq.NamedQubit("a")
b = cirq.NamedQubit("b")
gamma = 0.3 # Put your own value here.
circuit = cirq.Circuit.from_ops(cirq.ZZ(a,b)**gamma)
print(circuit)
cirq.unitary(circuit).round(2)

test_matrix = np.array([[np.exp(-1j*np.pi*gamma/2),0, 0, 0],
                        [0, np.exp(1j*np.pi*gamma/2), 0, 0],
                        [0, 0, np.exp(1j*np.pi*gamma/2), 0],
                        [0, 0, 0,np.exp(-1j*np.pi*gamma/2)]])
cirq.testing.assert_allclose_up_to_global_phase(test_matrix, cirq.unitary(circuit), atol=1e-5)

a = cirq.NamedQubit("a")
gamma = 0.3 # Put your own value here.
h = 1.3 # Put your own value here.
circuit = cirq.Circuit.from_ops(cirq.Z(a)**(gamma*h))
print(circuit)
print(cirq.unitary(circuit).round(2))

test_matrix = np.array([[np.exp(-1j*np.pi*gamma*h/2),0],
                        [0, np.exp(1j*np.pi*gamma*h/2)]])
cirq.testing.assert_allclose_up_to_global_phase(test_matrix, cirq.unitary(circuit), atol=1e-5)

n_cols = 3
n_rows = 3
h = 0.5*np.ones((n_rows,n_cols))

# Arranging the qubits in a list-of-lists like this makes them easy to refer to later.
qubits = [[cirq.GridQubit(i,j) for j in range(n_cols)] for i in range(n_rows)]


def beta_layer(beta):
  """Generator for U(beta, B) layer (mixing layer) of QAOA"""
  for row in qubits:
    for qubit in row:
      yield cirq.X(qubit)**beta
    
def gamma_layer(gamma, h):
  """Generator for U(gamma, C) layer of QAOA
  
  Args:
    gamma: Float variational parameter for the circuit
    h: Array of floats of external magnetic field values
  """
  for i in range(n_rows):
    for j in range(n_cols):
      if i < n_rows-1:
        yield cirq.ZZ(qubits[i][j], qubits[i+1][j])**gamma
      if j < n_cols-1:
        yield cirq.ZZ(qubits[i][j], qubits[i][j+1])**gamma
      yield cirq.Z(qubits[i][j])**(gamma*h[i,j])

qaoa = cirq.Circuit()
qaoa.append(cirq.H.on_each(*[q for row in qubits for q in row]))
# YOUR CODE HERE
print(qaoa)

qaoa = cirq.Circuit()
gamma = sympy.Symbol('g')
beta = sympy.Symbol('b')
qaoa.append(cirq.H.on_each(*[q for row in qubits for q in row]))
qaoa.append(gamma_layer(gamma,h))
qaoa.append(beta_layer(beta))
print(qaoa)

def energy_from_wavefunction(wf, h):
  """Computes the energy-per-site of the Ising Model directly from the
  a given wavefunction. 
  
  Args:
    wf: Array of size 2**(n_rows*n_cols) specifying the wavefunction.
    h: Array of shape (n_rows, n_cols) giving the magnetic field values.
  
  Returns:
    energy: Float equal to the expectation value of the energy per site 
  """
  
  n_sites = n_rows*n_cols

  # Z is an array of shape (n_sites, 2**n_sites). Each row consists of the 
  # 2**n_sites non-zero entries in the operator that is the Pauli-Z matrix on
  # one of the qubits times the identites on the other qubits. The
  # (i*n_cols + j)th row corresponds to qubit (i,j).
  Z = np.array([(-1)**(np.arange(2**n_sites) >> i) for i in range(n_sites-1,-1,-1)])

  # Create the operator corresponding to the interaction energy summed over all
  # nearest-neighbor pairs of qubits
  ZZ_filter = np.zeros_like(wf, dtype=float)
  for i in range(n_rows):
    for j in range(n_cols):
      if i < n_rows-1:
        ZZ_filter += Z[i*n_cols + j]*Z[(i+1)*n_cols + j]
      if j < n_cols-1:
        ZZ_filter += Z[i*n_cols + j]*Z[i*n_cols + (j+1)]

  energy_operator = -ZZ_filter - h.reshape(n_sites).dot(Z)
        
        
  # Expectation value of the energy divided by the number of sites
  return np.sum(np.abs(wf)**2 * energy_operator) / n_sites

def energy_from_params(gamma, beta, qaoa, h):
  sim = cirq.Simulator()
  params = cirq.ParamResolver({'g':gamma, 'b':beta})
  wf = sim.simulate(qaoa, param_resolver = params).final_state
  return energy_from_wavefunction(wf, h)  

%%time
grid_size = 50
gamma_max = 2
beta_max = 2

energies = np.zeros((grid_size,grid_size))
for i in range(grid_size):
  for j in range(grid_size):
    energies[i,j] = energy_from_params(i*gamma_max/grid_size, j*beta_max/grid_size, qaoa, h)

plt.ylabel('gamma')
plt.xlabel('beta')
plt.title('Energy as a Function of Parameters')
plt.imshow(energies, extent = (0,beta_max,gamma_max,0));

def gradient_energy(gamma, beta, qaoa, h):
  """Uses a symmetric difference to calulate the gradient."""
  eps = 10**-3 # Try different values of the discretization parameter
  
  # Gamma-component of the gradient
  grad_g = energy_from_params(gamma + eps, beta, qaoa, h)
  grad_g -= energy_from_params(gamma - eps, beta, qaoa, h)
  grad_g /= 2*eps
  
  # Beta-compoonent of the gradient
  grad_b = energy_from_params(gamma, beta + eps, qaoa, h)
  grad_b -= energy_from_params(gamma, beta - eps, qaoa, h)
  grad_b /= 2*eps  
  
  return grad_g, grad_b

gamma, beta = 0.2, 0.7 # Try different initializations
eta = 10**-2 # Try adjusting the learning rate.
for i in range(151):
  grad_g, grad_b = gradient_energy(gamma, beta, qaoa, h)
  gamma -= eta*grad_g
  beta -= eta*grad_b
  if not i%25:
    print('Step: {} Energy: {}'.format(i, energy_from_params(gamma, beta, qaoa, h)))
print('Learned gamma: {} Learned beta: {}'.format(gamma, beta, qaoa, h))

measurement_circuit = qaoa.copy()
measurement_circuit.append(cirq.measure(*[qubit for row in qubits for qubit in row],key='m'))
measurement_circuit

num_reps = 10**3 # Try different numbers of repetitions
gamma, beta = 0.2,0.25 # Try different values of the parameters
simulator = cirq.Simulator()
params = cirq.ParamResolver({'g':gamma, 'b':beta})
result = simulator.run(measurement_circuit, param_resolver = params, repetitions=num_reps)

def compute_energy(meas, h):
  Z_vals = 1-2*meas.reshape(n_rows,n_cols)
  energy = 0
  for i in range(n_rows):
    for j in range(n_cols):
      if i < n_rows-1:
        energy -= Z_vals[i, j]*Z_vals[i+1, j]
      if j < n_cols-1:
        energy -= Z_vals[i, j]*Z_vals[i, j+1]
      energy -= h[i,j]*Z_vals[i,j]
  return energy/(n_rows*n_cols)

hist = result.histogram(key='m')
num = 10
probs = [v/result.repetitions for _,v in hist.most_common(num)]
configs = [c for c,_ in hist.most_common(num)]

plt.title('Probability of {} Most Common Outputs'.format(num))
plt.bar([x for x in range(len(probs))],probs)
plt.show()
meas = [[int(s) for s in ''.join([str(b) for b in bin(k)[2:]]).zfill(n_rows*n_cols)] for k in configs]
costs = [compute_energy(np.array(m), h) for m in meas]
plt.title('Energy of {} Most Common Outputs'.format(num))
plt.bar([x for x in range(len(costs))],costs)
plt.show()
print('Fraction of outputs displayed: {}'.format(np.sum(probs).round(2)))