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

# # Defining and using custom Plugins
# 
# Inside the QLM, Plugins can be attached to QPUs in order to create more powerful computation stacks.
# 
# This notebook will present the Plugin interface and explain how to implement your own Plugin.
# 
# The **Plugin** interface is described inside a python abstract class `AbstractPlugin`:

# In[ ]:


from qat.core.plugins import AbstractPlugin


# 
# Let us build a simple Plugin that prints any circuit passing through on their way to the QPU and all results coming back to the user from the QPU.
# 
# 
# 

# In[ ]:


from qat.core import Batch, HardwareSpecs, Result
from qat.core.util import get_syntax

## As stated, our Plugin should inherit from the AbstractPlugin class:
class Printer(AbstractPlugin):
    ## Plugins must implement a compile method
    ## here batch will be a Batch object (see doc) that contains a list of Jobs
    ## Specs will contain the specs of the attached QPU
    def compile(self, batch : Batch, specs : HardwareSpecs) -> Batch:
        print("Compiling a batch of size", len(batch.jobs))
        for count, job in enumerate(batch.jobs):
            print("Job #", count)
            for index in range(len(job.circuit.ops)):
                syntax = get_syntax(job.circuit, index)
                print(syntax)
        ## We can use the meta_data field of the batch to store infos
        batch.meta_data = {"printer" : "this batch was printed!"}
        return batch
    
    ## Plugins must implement a post_process method
    ## here results will be a list of Results and
    ## meta_data a string,string dictionary that is attached to the corresponding
    ## Batch object
    def post_process(self, batch_result):
        print("We got", len(batch_result.results), "results")
        for result in batch_result.results:
            print(len(result), "samples in this result")
            for sample in result:
                print(sample)
        print("Sanity check : ", batch_result.meta_data) ## Just to check that we saw the Batch on its way in.
        return batch_result


# **Plugins** can be attached to **QPUs** using a pipe operator, or by invoquing the `push_plugin` method. If a job is submitted to the resulting QPU:
#  1. First, the circuit is compiled by each plugin (method `compile` of each plugin is called)
#  2. Second, the circuit is executed by the QPU
#  3. Then, results are post-processed by each plugin (method `post_process` of each plugin is called)

# In[ ]:


from qat.lang.AQASM import *
prog = Program()
qbits = prog.qalloc(2)
for qb in qbits:
    prog.apply(H, qb)


from qat.qpus import get_default_qpu
qpu = Printer() |Printer() |Printer() |Printer() |Printer() |Printer() | get_default_qpu()

job = prog.to_circ().to_job()
results = qpu.submit(job)


# If you are not convincedby the '|' notation, you can use a more traditional way of adding a plugin to a QPU:

# In[ ]:


other_qpu = get_default_qpu()
other_qpu.push_plugin(Printer())
results = other_qpu.submit(job)


# In[ ]: