"""
This module defines the Gateset class, which represents the allowed gates and topology for a specific quantum computer.

Several Implementations of Gateset are also defined here.
ZXZXZCNOTLinear -- A Gateset that uses CNOT and the ZXZXZ single qubit parameterization with the linear topology.
U3CNOTLinear -- A Gateset that uses CNOT and the U3 single qubit parameterization with the linear topology.
QubitCNOTLinear -- A Gateset that uses CNOT and the U3 single qubit parameterization with the linear topology, except it uses an XZXZ instead of a U3 after the control qubit of each CNOT.  This results in a gateset that covers the same search space as U3CNOTLinear, but with fewer redundant parameters, and therefore faster runtime.
QubitCNOTRing -- Uses U3 and XZXZ like QubitCNOTLinear, but includes a NonadjacentCNOTGate to add a link from the last qubit to the 0th.
QubitCNOTAdjacencyList -- Similar to QubitCNOTLinear and QubitCNOTRing, but takes in an adjacency list which uses NonadjacentCNOTGate to define work with a custom topology.
QutritCPIPhaseLinear -- A qutrit gateset that uses the CPIPhase gate as its two-qutrit gate, with a linear topology.
QutritCNOTLinear -- A qutrit gateset that uses an upgraded version of the CNOT gate as its two-qutrit gate, with a linear topology.

Several aliases are also defined, for the most common use cases.
DefaultQubit -- The default Gateset for working with qubits.  Currently is equivalent to QubitCNOTLinear.
DefaultQutrit -- The default Gateset for working with qutrits.  Currently is equivalent to QutritCPIPhaseLinear.
Default -- The overall default Gateset, which is equivalent to DefaultQubit.
"""
from .gates import *
from .assemblers import flatten_intermediate
import numpy as np



class Gateset():
    """This class defines the supported gates and topology for a specific quantum hardware."""
    def __init__(self):
        """Gatesets must set the value of d in their initializer, which represents the size of qudits that are supported (e.g. 2 for qubits or 3 for qutrits)."""
        self.d = 0
        raise NotImplementedError("Gatesets must implemented their own initializers and must set self.d to reflect the size of the qudits implemented in the gateset")

    def initial_layer(self, qudits):
        """
        The initial layer in the compilation.  Usually a layer of parameterized single-qudit gates.

        qudits -- The number of qudits in this circuit.
        expected return value -- A single Gate representing an initial layer for the circuit
        """
        return None # NOTE: Returns A SINGLE gate

    # the set of possible multi-qubit gates for searching.  Generally a two-qubit gate with single qubit gates after it.
    def search_layers(self, qudits):
        """
        A set of possible multi-qubit gates for searching.  Usually this is a two-qudit gate followed by two single-qudit gates, for every allowed placement of the two-qudit gate.  This defines the branching factor of the search tree.

        qudits -- The number of qudits in this circuit
        expected return value -- A list of tuples of (gate,weight) where Gate is the Gate representing that possible placement of the two-qudit gate, and weight is the weight or cost of adding that gate in that placement to the final circuit.
        """
        return [] # NOTES: Returns a LIST of tuples of (gate, weight)

    def branching_factor(self, qudits):
        """
        Returns an integer indicating the expected branching factor.  Usually this is automatically determined from search_layers, but it may need to be overridden if successors is overridden.

        qudits -- The number of qudits in this circuit
        expected return value -- An integer indicating the expecte branching factor
        """
        # This is the default implementation, for Gatesets that rely on search_layers
        return len(self.search_layers(qudits))

    def successors(self, circ, qudits=None):
        """
        Returns a list of Gates that are successors in the search tree to the input Gate, circ, representing a current ansatz circuit.

        circ -- The curret ansatz Gate.
        qudits -- The number of qudits in this circuit.
        expected return value -- A list of tuples of (gate, weight) where gate is a Gate that is a successor to circ, and weight is the cost or weight of moving to gate from circ.
        """

        # NOTE: Returns a LIST of tuples of (gate, weight)
        # NOTE: it is safe to assume that the circuit passed in here was produced by the functions of this class
        
        # This is the default implementation, for Gatesets that rely on search_layers
        qudits = int(np.log(circ.matrix([0]*circ.num_inputs).shape[0])/np.log(self.d))
        return [(circ.appending(t[0]), t[1]) for t in self.search_layers(qudits)]

    def __eq__(self, other):
        if self is other:
            return True
        if self.__module__ == Gateset.__module__:
            if type(self) == type(other):
                if self.__dict__ == other.__dict__:
                    return True
        return False

class ZXZXZCNOTLinear(Gateset):
    """A Gateset for working with CNOT and single-qubit gates parameterized with ZXZXZGate on the linear topology."""
    def __init__(self):
        self.single_gate = ZXZXZGate()
        self.cnot = CNOTGate()
        self.d = 2

    def initial_layer(self, n):
        return fill_row(self.single_gate, n)
    
    def search_layers(self, n):
        return linear_topology(self.cnot, self.single_gate, n, self.d)

class U3CNOTLinear(Gateset):
    """A Gateset for working with CNOT and single-qubit gates parameterized with U3Gate on the linear topology."""
    def __init__(self):
        self.single_gate = U3Gate()
        self.cnot = CNOTGate()
        self.d = 2

    def initial_layer(self, n):
        return fill_row(self.single_gate, n)
    
    def search_layers(self, n):
        return linear_topology(self.cnot, self.single_gate, n, self.d)

class QubitCNOTLinear(Gateset):
    """A Gateset for working with CNOT and single-qubit gates parameterized with U3Gate and XZXZGate on the linear topology.  This Gateset covers the same search space but uses fewer parameters than ZXZXZCNOTLinear and U3CNOTLinear."""
    def __init__(self):
        self.single_gate = U3Gate()
        self.single_alt  = XZXZGate()
        self.cnot = CNOTGate()
        self.d = 2

    def initial_layer(self, n):
        return fill_row(self.single_gate, n)
    
    def search_layers(self, n):
        return linear_topology(self.cnot, self.single_gate, n, self.d, single_alt=self.single_alt)

    def branching_factor(self, qudits):
        return qudits-1

    def successors(self, circ, qudits=None):
        if qudits is None:
            qudits = int(np.log(circ.matrix([0]*circ.num_inputs).shape[0])/np.log(self.d))
        skip_index = find_last_3_cnots_linear(circ)
        return [(circ.appending(layer[0]), layer[1]) for layer in linear_topology(self.cnot, self.single_gate, qudits, self.d, single_alt=self.single_alt, skip_index=skip_index)]

class QubitCNOTRing(Gateset):
    """A Gateset for working with CNOT and single-qubit gates parameterized with U3Gate and XZXZGate on the ring topology."""
    def __init__(self):
        self.single_gate = U3Gate()
        self.single_alt  = XZXZGate()
        self.cnot = CNOTGate()
        self.d = 2

    def initial_layer(self, n):
        return fill_row(self.single_gate, n)

    def search_layers(self, n):
        I = IdentityGate()
        gates = linear_topology(self.cnot, self.single_gate, n, self.d, identity_gate=I, single_alt=self.single_alt)
        if n == 2:
            return gates
        finisher = (ProductGate(NonadjacentCNOTGate(n, n-1, 0), KroneckerGate(self.single_gate, *[I]*(n-2), self.single_alt)), 1)
        return gates + [finisher]


class QubitCNOTAdjacencyList(Gateset):
    """A Gateset for working with CNOT and single-qubit gates parameterized with U3Gate and XZXZGate on a custom topology, specified in the initializer."""
    def __init__(self, adjacency):
        """
        Allows the specification of a custom topology through an adjacency list.

        adjacency -- A list of tuples specifying which CNOT placements are allowed.  The tuples must be in the form of (control, target).

        For example, this is how you would specifiy the ring topology for 3 qubits:
        [(0,1), (1,2), (2,1)]

        It is not recommended to add bi-directional links, because with the arbitrary parameterized single qubit gates everywhere, such links would be redundant.
        """
        self.single_gate = U3Gate()
        self.single_alt = XZXZGate()
        self.I = IdentityGate()
        self.adjacency = adjacency
        self.d = 2

    def initial_layer(self, n):
        return fill_row(self.single_gate, n)

    def search_layers(self, n):
        if n == 2:
            return [(ProductGate(CNOTGate(), KroneckerGate(self.single_gate, self.single_gate)), 1)] # prevents the creation of an extra cnot placement in the 2 qubit case

        gates = []
        for pair in self.adjacency:
            if pair[0] >= n or pair[1] >= n:
                continue
            cnot = NonadjacentCNOTGate(n, pair[0], pair[1])
            single_gates = [self.single_gate if i == pair[0] else self.single_alt if i == pair[1] else self.I for i in range(0, n)]
            gates.append((ProductGate(cnot, KroneckerGate(*single_gates)), 1)) 
        return gates

class QutritCPIPhaseLinear(Gateset):
    """A Gateset for working with CPIPhase and single-qutrit gates on the linear topology."""
    def __init__(self):
        self.single_gate = SingleQutritGate()
        self.d = 3

    def initial_layer(self, n):
        return fill_row(self.single_gate, n)
    
    def search_layers(self, n):
        return linear_topology(CPIPhaseGate(), self.single_gate, n, self.d)

class QutritCNOTLinear(Gateset):
    """A hybrid Gateset for working with CNOT and single-qutrit gates on the linear topology."""
    def __init__(self):
        self.single_gate = SingleQutritGate()
        self.d = 3

    def initial_layer(self, n):
        return fill_row(self.single_gate, n)
    
    def search_layers(self, n):
        return linear_topology(UpgradedConstantGate(CNOTGate()), self.single_gate, n, self.d)

# Commonly used functions for generating gatesets
def linear_topology(double_gate, single_gate, n, d, identity_gate=None, single_alt=None, double_weight=1, single_weight=0, skip_index=None):
    weight = double_weight + 2*single_weight
    if not identity_gate:
        identity_gate = IdentityGate(d=d)
    if single_alt is None:
        return [(KroneckerGate(*[identity_gate]*i, ProductGate(double_gate, KroneckerGate(single_gate, single_gate)), *[identity_gate]*(n-i-2)), weight) for i in range(0, n-1) if skip_index is None or i != skip_index]
    else:
        return [(KroneckerGate(*[identity_gate]*i, ProductGate(double_gate, KroneckerGate(single_alt, single_gate)), *[identity_gate]*(n-i-2)), weight) for i in range(0, n-1) if skip_index is None or i != skip_index]

def fill_row(gate, n):
    return KroneckerGate(*[gate]*n)


def find_last_3_cnots_linear(circuit):
    # this function finds the last 3 CNOTs in the circuit and returns the index if they are all in a row or None otherwise
    # this function is written specifically for linear topology-based gatesets.  Other gatesets should use find_last_3_cnots_arbitrary.
    # do not use either of these functions with gatesets not intended specifically for it

    contents = circuit._subgates
    if len(contents) < 4:
        return None
    indices = [[type(sub) for sub in gate._subgates].index(ProductGate) for gate in contents[-1:-4:-1]]

    if indices[0] == indices[1] and indices[1] == indices[2]:
        return indices[0]
    else:
        return None


# commonly used defaults
DefaultQubit = QubitCNOTLinear
DefaultQutrit = QutritCPIPhaseLinear
Default = DefaultQubit

