#!/usr/bin/env python
# coding: utf-8

# # Grover and quantum search
# 
# The purpose of this notebook is to briefly demonstrate the programming and execution of a simple Quantum search algorithm in the QLM.
# 
# ## Grover's algorithm
# 
# Grover's algorithm rely on a two main ingredients:
# - a diffusion operator $\mathcal{D} = 2 |s\rangle\langle s| - I$, where $|s\rangle = \frac{1}{\sqrt{2^n}}\sum |i\rangle$
# - and an oracle "marking" some basis states by flipping their phases: $O_f = \sum_i (-1)^{f(i)}|i\rangle\langle i|$ for some boolean predicate $f$
# 
# 
# ### Diffusion operator
# 
# Lets first start by programming a piece of pyAQASM code that will generate a diffusion operator.
# $\mathcal{D}$ can be obtained by:
# - performing a basis change from the computation basis to the diagonal basis (i.e a cascade of H gates)
# - performing a collectiong of bit flips in order to send the $|0^n\rangle$ state onto $|1^n\rangle$
# - performing a "global" controlled-Z gate in order to flip the phase of $|0^n\rangle
# - undoing the bit flips
# - undoing the basis change
# 
# 
# This can be easily expressed in pyAQASM as follows:

# In[ ]:


from qat.lang.AQASM import QRoutine, Z, H, X, CNOT
from qat.lang.AQASM.misc import build_gate

@build_gate("D", [int], arity=lambda n: n)
def diffusion(nbqbits):
    rout = QRoutine()
    wires = rout.new_wires(nbqbits)
    with rout.compute():
        for wire in wires:
            H(wire)
            X(wire)
    Z.ctrl(nbqbits - 1)(wires)
    rout.uncompute()
    return rout

example = diffusion(4)
example.display()
example.display(depth=1)


# Now we have a function that produces diffusion operators of the correct size. Moreover, this function is wrapped in a "box" called "D" for diffusion (this is mostly for visualization purpose).
# 
# ### Oracle
# 
# In order to produce a self-contained example, we will assume that we want to run a Grover instance for a particular predicate characterizing clean colorations of a graph $G$.
# 
# Given some graph $G(V, E)$, a clean coloration is a function $c: V \rightarrow \mathcal{C}$ for some finite set $\mathcal{C}$ such that for all edge $\{i, j\}\in E$, we have $c(i) \neq c(j)$.
# 
# #### Representing colorations:
# First, we need to represent graph coloration using only qubits. To do so, we will consider that $\mathcal{C}= [0, ...,k -1]$ and use $log(k)$ qubits to represent each color.
# To keep things simple, we will assume that the number of colors is a power of two: $k = 2^m$. That way, we will use $m$ qubits to store the color of each vertex.
# 
# This brings up the number of qubits to $n.m$ where $n = |V|$.
# 
# #### Checking colorations
# 
# In order to check that a coloration is clean, we need to check that for each edge $\{i, j\}$ we have $c(i) \neq c(j)$.
# This can be easily checked by:
# - xoring the two colors $c(i)$ and $c(j)$
# - checking that the result is not $|0\rangle$
# - computing the logical AND of all these conditions
# 
# Lets first write a small routine that will check a single edge:

# In[ ]:


@build_gate("CHECK", [int], arity=lambda m:2 * m + 1)
def check_edge(m):
    rout = QRoutine()
    color1 = rout.new_wires(m)
    color2 = rout.new_wires(m)
    output = rout.new_wires(1)
    with rout.compute():
        for wire1, wire2 in zip(color1, color2):
            CNOT(wire1, wire2)
        for wire in color2:
            X(wire)
    X.ctrl(m)(color2, output)
    rout.uncompute()
    X(output)
    return rout

example = check_edge(3)
example.display()
example.display(depth=1)


# Now that we have a routine checking color for a single edge, we can easily write a routine check the coloration of a full graph.
# 
# We need to compute a large AND of $|E|$ clauses. This can be done using $|E|$ ancillae and a single, large Toffoli gate with $|E|$ controls.
# This works, but is quite expensive. In order to slightly optimize the number of qubits, we will use the following trick:
# - initialize some quantum register to $L = |0\rangle$ of size $log |E|$
# - for each edge, if the corresponding clause is verified, increment $L$ by one
# - finally check that $L = |E|$ and flip the phase of the classical state. If $L \neq |E|$ , some clause was violated and the coloration is not clean!

# In[ ]:


import networkx as nx
# we need an adder:
from qat.lang.AQASM.qftarith import add

@build_gate("CHECK_GRAPH", [nx.Graph, int], arity=lambda g, m: g.number_of_nodes() * m)
def check_graph(graph, m):
    rout = QRoutine()
    # Colors array
    colors = [rout.new_wires(m) for _ in graph.nodes()]
    # Our counter L
    size_l = graph.number_of_edges().bit_length()
    L = rout.new_wires(size_l)
    # a work qubit to store the result of a check_edge
    tmp = rout.new_wires(1)
    # some routines (check_edge and an adder)
    check_routine = check_edge(m)
    adder = add(size_l, 1)
    with rout.compute():
        for a, b in graph.edges():
            # checking the edge
            with rout.compute():
                check_routine(colors[a], colors[b], tmp)
            # adding the result in our counter
            adder(L, tmp)
            # uncomputing 'tmp'
            rout.uncompute()
    # checking if l = |E|
    E = graph.number_of_edges()
    with rout.compute():
        for i in range(size_l):
            if ((E >> i) & 1) == 0:
                X(L[i])
    Z.ctrl(size_l - 1)(L)
    # uncomputing the X's
    rout.uncompute()
    # uncomputing L
    rout.uncompute()
    # tmp and L are work qubits and should be flagged to that they can be re-used
    rout.set_ancillae(L, tmp)
    return rout
    
graph = nx.generators.path_graph(3)
example = check_graph(graph, 2)

example.display()
example.display(depth=2)


# Ok, we're almost there! We now have our oracle that inverts the phase of a classical state if and only if it describes a clean coloring.
# We can now write our full Grover algorithm:

# In[ ]:


import numpy as np
from qat.lang.AQASM import Program
from qat.lang.AQASM.qint import QInt
from qat.qpus import get_default_qpu

def graph_coloration_grover(graph, m, nbsteps):
    prog = Program()
    # Notice that each color is stored in a different register with type "int"
    # This is so it is easy to display final measurement results in the next cells
    data = [prog.qalloc(m, QInt) for _ in range(graph.number_of_nodes())]
    # Initializing a superposition of all possible coloring
    for color in data:
        for qbit in color:
            H(qbit)
    # Our diffusion routine and our oracle
    O = check_graph(graph, m)
    D = diffusion(graph.number_of_nodes() * m)
    # The main loop
    for _ in range(nbsteps):
        O(data)
        D(data)
    # We also return the list of registers carying the colors to simplify data retrieval
    return prog.to_circ(), data


# Lets run a Grover algorithm on our path graph with 4 possible colors per nodes (completely overkill) and for 2 iterations.
# 
# We will iterate over the possible measurement outcomes and compute the success probability of the algorithm:

# In[ ]:


color, reg = graph_coloration_grover(graph, 2, 2)
print("Using", color.nbqbits, "qubits")
qpu = get_default_qpu()
# Here we specify that we want to measure the color registers
# the results will be formatted nicely thanks to this (using the .value field below)
result = qpu.submit(color.to_job(qubits=reg))

success_proba = 0
for sample in result:
    c1, c2, c3 = sample.state.value
    if c1 != c2 != c3:
        success_proba += sample.probability
print("Overall success probability:", success_proba)


# It seems to work quite alright. Lets try with a larger graph. 
# We will:
# - generate a random graph
# - run a grover to look for clean 4-colorations for several number of iterations
# - print the success probability for each number of iterations
# 
# The main reason why we need this for loop is because we can't predict the proportion of correct colorations among the full set of colorations.
# 
# In practice, the correct way of tackling this issue would be to use a slightly more advanced quantum search algorithm by Brassard et al.

# In[ ]:


graph = nx.generators.erdos_renyi_graph(5, 0.5)

color, reg = graph_coloration_grover(graph, 2, 0)
print("Using", color.nbqbits, "qubits")
qpu = get_default_qpu()

for nbsteps in range(2, 8):
    result = qpu.submit(color.to_job(qubits=reg))
    color, reg = graph_coloration_grover(graph, 2, nbsteps)
    success_proba = 0
    for sample in result:
        colors = sample.state.value
        if all(colors[a] != colors[b] for a, b in graph.edges()):
            success_proba += sample.probability
    print(nbsteps, success_proba)


# # Going further: QSearch
# 
# As mentionned before, the main issue of our current algorithm is that we can't predict the optimal number of iterations required to find a clean coloration with high probability.
# 
# In 2000, Brassard et al proposed a solution for that particular problem (which is quite recurrent when dealing with quantum search).
# All the details can be found [here](https://arxiv.org/abs/quant-ph/0005055).
# 
# The algorithm is a classical loop consisting of many calls to a Grover search with number of steps growing exponentially.
# 
# For simplicity, lets first define a wrapping function that will:
# - build a Grover search for a given number of steps
# - sample a single bitstring out of the final quantum state and return the corresponding coloration

# In[ ]:


def single_search(graph, m, nbsteps):
    grover, registers = graph_coloration_grover(graph, m, nbsteps)
    job = grover.to_job(qubits=registers, nbshots=1)
    result = qpu.submit(job)
    return result[0].state.value    


# We can now declare a `q_search` function implementing the classical loop as in the paper. For clarity we will try to use the same notations as the one in the reference.

# In[ ]:


def q_search(graph, m):
    l = 1
    c = 1.5
    M = int(np.ceil(c ** l))
    while True:
        j = np.random.randint(1, M)
        print("Trying with", j, "steps")
        coloration = single_search(graph, m, j)
        if all(coloration[a] != coloration[b] for a, b in graph.edges()):
            return coloration
        l += 1
        M = int(np.ceil(c ** l))


# In[ ]:


coloration = q_search(graph, 2)
print("Found a valid coloration:", coloration)


# Of course all this is just an brief introduction to quantum search algorithms!
# There are plenty of alternatives and variants that would deserve hundreds of notebooks.
#