import pyquil.forest as qvm_module
import numpy as np
from grove.pyqaoa.maxcut_qaoa import maxcut_qaoa
barbell = [(0, 1)]  # graph is a defined by a list of edges.  Edge weights are assumed to be 1.0
steps = 1  # evolution path length between the ref hamiltonian and cost hamiltonian
inst = maxcut_qaoa(barbell, steps=steps)  # initializing problem instance

"""
Explanation: The Quantum Approximate Optimization Algorithm for MAX-CUT
The following is a step-by-step guide to running QAOA on the MaxCut problem.  In the debue paper on QAOA (arXiv: 1411.4028), Farhi, Goldstone, and Gutmann demonstrate that the lowest order approximation of the algorithm produced an approximation ratio of 0.6946 for the MaxCut problem on three-regular graphs.  You can use this notebook to set up an arbitary graph for MaxCut and solve it using the QAOA algorithm the Rigetti Forest service.
pyQAOA is a python library that implements the QAOA.  It uses the PauliTerm and PauliSum objects from the pyQuil library for expressing the cost Hamiltonian and driver Hamiltonain. These operators are used to create a parametric pyQuil program and passed to the variational quantum eigensolver (VQE) solver in Grove.  VQE calls the Rigetti Forest QVM to exectue the Quil program that prepares the angle parameterized state.  There are multiple ways to construct the MAX-CUT problem for the QAOA library.  We include a method that accepts a graph and returns a QAOA instance where the costs and driver Hamiltonaians have been constructed.  The graph is either an undirected Networkx graph or a list of tuples where each tuple represents an edge between a pair of nodes.
We start by demonstrating the QAOA algorithm with the simplest instance of MAX-CUT--parititioning the nodes on a barbell graph.  The barbell graph corresponds to a single edge connecting two nodes.  The solution is the partitioning of the nodes into different sets ${0, 1}$.
End of explanation
"""


cost_list, ref_list = inst.cost_ham, inst.ref_ham
cost_ham = reduce(lambda x,y: x + y, cost_list)
ref_ham = reduce(lambda x,y: x + y, ref_list)
print cost_ham
print ref_ham

"""
Explanation: The cost Hamiltonian and driver Hamiltonian corresponding to the barbell graph are stored in QAOA object fields in the form of lists of PauliSums.
End of explanation
"""


param_prog = inst.get_parameterized_program()
prog = param_prog([1.2, 4.2])
print prog

"""
Explanation: The identity term above is not necessary to the computation since global phase rotations on the wavefunction do not change the expectation value.  We include it here purely as a demonstration.  The cost function printed above is the negative of the traditional maximum cut cost operator.  This is because QAOA is formulated as the maximization of the cost operator but the VQE algorithm in the pyQuil library performs a minimization.
QAOA requires the construction of a state parameterized by $\beta$ and $\gamma$ rotation angles
$$\begin{align}
\mid \beta, \gamma \rangle = \prod_{p=0}^{\mathrm{steps}}\left( U(\hat{H}{\mathrm{drive}}, \beta{p})U(\hat{H}{\mathrm{MAXCUT}}, \gamma{p}) \right)^{\mathrm{steps}} (\mid +\rangle_{N-1}\otimes\mid + \rangle_{N-2}...\otimes\mid + \rangle_{0}).
\end{align}$$
The unitaries $U(\hat{H}{\mathrm{drive}}, \beta{p})$ and $U(\hat{H}{\mathrm{MAXCUT}}, \gamma{p})$ are the exponentiation of the driver Hamiltonian and the cost Hamiltonian, respectively. 
$$
\begin{align}
U(\hat{H}{\mathrm{ref}}, \beta{p}) = e^{-i \beta_{p} \hat{H}{drive}} \
U(\hat{H}{\mathrm{MAXCUT}}, \gamma_{p}) = e^{-i \gamma_{p} \hat{H}_{\mathrm{MAXCUT}}}
\end{align}
$$
The QAOA algorithm relies on many constructions of a wavefunction via parameterized Quil and measurements on all qubits to evaluate an expectation value.  In order avoid needless classical computation, QAOA constructs this parametric program once at the beginning of the calculation and then uses this same program object throughout the computation.  This is accomplished using the ParametricProgram object from pyQuil that allows us to slot in a symbolic value for a parameterized gate.  
The parameterized program object can be accessed through the QAOA method get_parameterized_program().  Calling this method on an instantiated QAOA object returns a closure with a precomputed set of Quil programs.  Calling this closure with the parameters $\beta$ and $\gamma$ returns the circuit that has the rotations parameterized.
End of explanation
"""


betas, gammas = inst.get_angles()
print betas, gammas

"""
Explanation: The printed program above is a Quil program that can be executed on a QVM.  QAOA has two modes of operation: 1) pre-computing the angles of rotation classically and using the quantum computer to measure expectation values through repeated experiments and 2) installing the a classical optimization loop on top of step 1 to optimally determine the angles.  Operation mode 2 is known as the variational quantum eigensolver algorithm.  the QAOA object wraps the instantiation of the VQE algorithm with the get_angles() method.
End of explanation
"""


param_prog = inst.get_parameterized_program()
t = np.hstack((betas, gammas))
prog = param_prog(t)
wf, _ = inst.qvm.wavefunction(prog)
wf = wf.amplitudes
for ii in range(2**inst.n_qubits):
    print inst.states[ii], np.conj(wf[ii])*wf[ii]

"""
Explanation: get_angles() returns optimal the $\beta$ and $\gamma$ angles.  To view the probabilities of the state you can call QAOA.probabilities(t) were t is a concatentation of the $\beta$ and $\gamma$ angles, in that order.  The probabilities(t) routine takes the $\beta$ and $\gamma$ parameters, reconstructs the wave function and returns their coefficients.  A modified version can be used to print off the probabilities
End of explanation
"""


%matplotlib inline

from grove.pyqaoa.qaoa import QAOA
import networkx as nx
import matplotlib.pyplot as plt
from pyquil.paulis import PauliSum, PauliTerm
import pyquil.quil as pq
from pyquil.gates import H
import pyquil.forest as qvm_module
CXN = qvm_module.Connection()

# define a 6-qubit ring
ring_size = 6
graph = nx.Graph()
for i in range(ring_size):
    graph.add_edge(i, (i + 1) % ring_size)

nx.draw_circular(graph, node_color="#6CAFB7")

"""
Explanation: As expected the bipartitioning of a graph with a single edge connecting two nodes corresponds to the state ${|01\rangle, |10\rangle }$.  In this trivial example the QAOA finds angles that construct a distribution peaked around the two degenerate solutions.
MAXCUT on larger graphs and alternative optimizers.
Larger graph instances and different classical optimizers can be used with the QAOA.  Here we consider an 6-node ring of disagrees.  For an even number ring graph, the ring of disagrees corresponds to the antiferromagnet ground state--i.e. alternating spin-up spin-down.
End of explanation
"""


cost_operators = []
driver_operators = []
for i, j in graph.edges():
    cost_operators.append(PauliTerm("Z", i, 0.5)*PauliTerm("Z", j) + PauliTerm("I", 0, -0.5))
for i in graph.nodes():
    driver_operators.append(PauliSum([PauliTerm("X", i, 1.0)]))

"""
Explanation: This graph could be passed to the maxcut_qaoa method and a QAOA instance with the correct driver and cost Hamiltonian could be generated as before.  In order to demonstrate the more general approach, along with some VQE options, we will construct the cost and driver Hamiltonians directly with PauliSum and PauliTerm objects.  To do this we parse the edges and nodes of the graph to construct the relevant operators.
$$
\begin{align}
\hat{H}{\mathrm{cost}} = \sum{\langle i, j\rangle \in E}\frac{\sigma_{i}^{z}\sigma_{j}^{z} - 1}{2} \
\hat{H}{\mathrm{drive}} = \sum{i}^{n}-\sigma_{i}^{x}
\end{align}
$$
where $\langle i, j\rangle \in E$ refers to the pairs of nodes that form the edges of the graph.
End of explanation
"""


prog = pq.Program()
for i in graph.nodes():
    prog.inst(H(i))

"""
Explanation: We will also construct the initial state and pass this to the QAOA object.  By default QAOA uses the $|+\rangle$ tensor product state.  In other notebooks we will demonstrate that you can use the driver_ref optional argument to pass a different starting state for QAOA.
End of explanation
"""


ring_cut_inst = QAOA(CXN, len(graph.nodes()), steps=1, ref_ham=driver_operators, cost_ham=cost_operators,
                     driver_ref=prog, store_basis=True, rand_seed=42)

betas, gammas = ring_cut_inst.get_angles()

"""
Explanation: Now we are ready to instantiate the QAOA object!
End of explanation
"""


from collections import Counter

# get the parameterized program
param_prog = ring_cut_inst.get_parameterized_program()
sampling_prog = param_prog(np.hstack((betas, gammas)))

# use the run_and_measure QVM API to prepare a circuit and then measure on the qubits
bitstring_samples = CXN.run_and_measure(sampling_prog, range(len(graph.nodes())), 1000)
bitstring_tuples = map(tuple, bitstring_samples)

# aggregate the statistics
freq = Counter(bitstring_tuples)
most_frequent_bit_string = max(freq, key=lambda x: freq[x])
print freq

print "The most frequently sampled string is ", most_frequent_bit_string

"""
Explanation: We are interested in the bit strings returned from the QAOA algorthm.  The get_angles() routine calls the VQE algorithm to find the best angles.  We can then manually query the bit strings by rerunning the program and sampling many outputs.
End of explanation
"""


# plotting strings!
n_qubits = len(graph.nodes())
def plot(inst, probs):
    probs = probs.real
    states = inst.states
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.set_xlabel("state",fontsize=20)
    ax.set_ylabel("Probability",fontsize=20)
    ax.set_xlim([0, 2**n_qubits ])
    rec = ax.bar(range(2**n_qubits), probs, )
    num_states = [0, int("".join(str(x) for x in [0,1] * (n_qubits/2)), 2), 
              int("".join(str(x) for x in [1,0] * (n_qubits/2)), 2), 2**n_qubits - 1 ]
    ax.set_xticks(num_states)
    ax.set_xticklabels(map(lambda x: inst.states[x], num_states), rotation=90)
    plt.grid(True)
    plt.tight_layout()
    plt.show()

t = np.hstack((betas, gammas))
probs = ring_cut_inst.probabilities(t)
plot(ring_cut_inst, probs)

"""
Explanation: We can see that the first two most frequently sampled bit strings are the alternating solutions to the ring graph.  Since we have access to the wave function we can go one step farther and view the probability distrubtion over the bit strings produced by our $p = 1$ circuit.
End of explanation
"""


# get the angles from the last run
beta = ring_cut_inst.betas
gamma = ring_cut_inst.gammas
# form new beta/gamma angles from the old angles
betas = np.hstack((beta[0]/3, beta[0]*2/3))
gammas = np.hstack((gamma[0]/3, gamma[0]*2/3))
# set up a new QAOA instance.
ring_cut_inst_2 = QAOA(CXN, len(graph.nodes()), steps=2, ref_ham=driver_operators, cost_ham=cost_operators,
                     driver_ref=prog, store_basis=True, init_betas=betas, init_gammas=gammas)
# run VQE to determine the optimal angles
betas, gammas = ring_cut_inst_2.get_angles()
t = np.hstack((betas, gammas))
probs = ring_cut_inst_2.probabilities(t)
plot(ring_cut_inst_2, probs)

"""
Explanation: For larger graphs the probability of sampling the correct string could be significantly smaller, though still peaked around the solution.  Therefore, we would want to increase the probability of sampling the solution relative to any other string.  To do this we simply increase the number of steps $p$ in the algorithm.  We might want to bootstrap the algorithm with angles from lower number of steps.  We can pass inital angles to the solver as optional arguments.
End of explanation
"""


from scipy.optimize import fmin_bfgs

ring_cut_inst_3 = QAOA(CXN, len(graph.nodes()), steps=3, ref_ham=driver_operators, cost_ham=cost_operators,
                       driver_ref=prog, store_basis=True, minimizer=fmin_bfgs, minimizer_kwargs={'gtol':1.0e-3},
                       rand_seed=42)
betas, gammas = ring_cut_inst_3.get_angles()
t = np.hstack((betas, gammas))
probs = ring_cut_inst_3.probabilities(t)
plot(ring_cut_inst_3, probs)

"""
Explanation: We could also change the optimizer which is passed down to VQE through the QAOA interface.  Let's say I want to use BFGS or another optimizer that can be wrapped in python.  Simply pass it to QAOA through the minimzer, minimizer_args, and minimizer_kwargs keywords
End of explanation
"""


