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

# # Very advanced usage: linking under control
# 
# In this notebook, we present how to further exploit the power of the pyAQASM linking mechanics.
# We assume that you are already familiar with the basic linking mechanic and ancillae management.
# 
# ## A story of Toffolis
# 
# The example we will present here is related to a quite common situation when dealing with classical oracles and Grover algorithm.
# 
# Lets assume that we would like to implement a Grover diffusion operator over an arbitrary number of bits.
# 
# Grover's diffusion operator has shape:
# 
# $$ U_s = 2|s\rangle\langle s| - I $$
# where $|s\rangle = |+^n\rangle$ is a uniform superposition of all possible classical states over $n$ qubits.
# 
# It is easy to implement this transformation using a single $n$-qubits Toffoli gate (i.e with $n-1$ controls) as folows:

# In[ ]:


from qat.lang.AQASM import QRoutine, H, CCNOT, X

def grover_diffusion(nbqbits):
    diffusion = QRoutine()
    wires = diffusion.new_wires(nbqbits)
    with diffusion.compute():
        for wire in wires:
            H(wire)
            X(wire)
        H(wires[-1])
    CCNOT.ctrl(nbqbits - 3)(wires)
    diffusion.uncompute()
    return diffusion
diffusion_5 = grover_diffusion(5)
diffusion_5.display()


# So we now have a routine that flawlessly implement a Grover's diffusion.
# 
# In a real application however, we do not have access to such a large gate. Usually, a generalized Toffoli is implemented using ancillae qubits a many smaller gates (CNOT + T or 3-qubits Toffoli gates).
# 
# In our current situation, we already used the CCNOT gate, controlled many times, to implement our routine. What could we do to postpone the choice of implementation of the multi-controlled Toffoli to after having written the full algorithm?
# 
# This is what we are going to demonstrate here.
# 
# ## Linking under a (or many) control(s)
# 
# 
# The linking process of a sub-circuit implementation to a gate roughly works like this:
#  * the linker crawls the circuit and find gate that have no subcircuit implementation but whose definition has a circuit generator
#  * for each of these gates, the linker :
#      * (i) pops all operators such as ctrls and daggers off of the gate
#      * (ii) calls the circuit generator of the gate to produce a subcircuit (i.e a QRoutine object)
#      * (iii) re-applies the operators on the QRoutine
#      * (iv) inserts this subcircuit in the definition of the gate
# 
# Hence, if we manager to describe our Toffoli gate as some kind of QRoutine that has a particular behavior when controlled, then we can finely control the way our multi-control Toffoli will be generated in the final circuit.
# 
# We start by creating a new class inheriting from QRoutine, with a single overloaded method:

# In[ ]:


class MultiToffoli(QRoutine):
    def __init__(self):
        super().__init__()
        # Since we want to modify the definition of CCNOT, we better not use CCNOT inside here
        # If we were to use a CCNOT here, the linker would loop indefinitly
        X.ctrl(2)(self.new_wires(3)) 
        
    def ctrl(self, nbctrls=1):
        if nbctrls == 0:
            # Same argument here
            return X.ctrl(2)
        rout = QRoutine()
        wires = rout.new_wires(self.arity + nbctrls)
        tmp = rout.new_wires(1)
        rout.set_ancillae(tmp)
        with rout.compute():
            # And here
            X.ctrl(2)(wires[0], wires[1], tmp)
        self.ctrl(nbctrls - 1)(tmp, wires[2:])
        rout.uncompute()
        return rout

m_ccnot = MultiToffoli().ctrl(1)
m_ccnot.display()
m_ccnot = m_ccnot.ctrl(3)
m_ccnot.display()


# So now we have a QRoutine that knows how to better deal with control operators, instead of naively increasing its arity.
# Notice that this implementation if far from being "smart" and can be optimized quite a lot! But this is not the point of this notebook.
# 
# Now, we would like to link this implementation of the CCNOT in a Program that was generated using our previous definition of the Grover operator.
# 
# To do that we need to define a new CCNOT gate that contains the as circuit implementation the class we just defined:

# In[ ]:


from qat.lang.AQASM.misc import build_gate

@build_gate("CCNOT", [], arity=3)
def multi_ccnot():
    return MultiToffoli()


# And we can link it to our program:

# In[ ]:


from qat.lang.AQASM import Program
my_algorithm = Program()
qbits = my_algorithm.qalloc(6)
grover_diffusion(6)(qbits)
circuit1 = my_algorithm.to_circ(include_matrices=False) # Without linking

circuit2 = my_algorithm.to_circ(link=[multi_ccnot], include_matrices=False) # With linking
circuit1.display()
circuit2.display()
circuit2.display(depth=1)


# ## Linking after circuit generation
# 
# Lets go a bit further, and assume that a circuit was already generated without linking our own implementation of the multi-Toffoli.

# In[ ]:


circuit = my_algorithm.to_circ(include_matrices=False)
circuit.display(depth=1)
from qat.lang.linking.linker import Linker
import qat.lang.linking.util as linkutil
linker = Linker(link=[multi_ccnot], include_matrices=False)
# This replaces the multi-control Toffolis by their proper implementation
linker.link(circuit)
circuit.display(depth=1)


# And, by the way, if you are lost among the qubits, keeping the lock/release of ancillae visible can help debug:

# In[ ]:


circuit = my_algorithm.to_circ(include_matrices=False)
from qat.lang.linking.linker import Linker
import qat.lang.linking.util as linkutil
linker = Linker(link=[multi_ccnot], include_matrices=False, include_locks=True)
# This replaces the multi-control Toffolis by their proper implementation
linker.link(circuit)
circuit.display(depth=1)


# In[ ]: