import random
import numpy as np

class Gate():
    def __init__(self, childGates=[]):
        self.childGates = childGates
        self._condense()

    def _condense(self):
        self.mat = np.identity(2)
        for gate in self.childGates:
            self.mat = np.dot(self.mat, gate.mat)
        return True

    def __call__(self,*arg):
        if len(arg)==0:
            return self
        elif len(arg)==1:
            return self._mul(inp)
        else:
            raise Exception("NaBrO")

    def __mul__(self, other):
        return self._mul(other)

    def __repr__(self):
        return "<qGate ["+str(self)+"]>"

    def __str__(self):
        return "".join(str(gate) for gate in self.childGates)

    def _mul(self, obj):
        if isinstance(obj, Gate):
            return self._op(obj)
        elif isinstance(obj, State):
            return State(np.dot(self.mat, obj.vec))

    def _op(self, obj):
        return Gate(self.childGates + [obj])

    @property
    def adj(self):
        return AdjGate(self)

    def shorten(self):
        newGates = []
        changes = 0
        i=0
        while i < len(self.childGates):
            gate = self.childGates[i]
            if i < len(self.childGates)-1:
                nextChild = self.childGates[i+1]
            else:
                nextChild = None
            changes += 1
            if np.all(gate.mat==np.identity(2)):
                pass
            elif isinstance(gate, AdjGate) and isinstance(gate.origGate, AdjGate):
                newGates.append(gate.origGate.origGate)
            elif isinstance(gate, AdjGate) and str(gate.origGate) == "H":
                newGates.append(H)
            elif isinstance(gate, AdjGate) and str(gate.origGate) == str(nextChild):
                newGates.append(gate.origGate)
                i += 1
            elif i != len(self.childGates) and isinstance(nextChild, AdjGate) and str(nextChild.origGate) == str(gate):
                newGates.append(nextChild.origGate)
                i += 1
            else:
                newGates.append(gate)
                changes -= 1
            i += 1

        if len(newGates)==0:
            newGates = I()

        newGate = Gate(newGates)

        if changes != 0:
            return newGate.shorten()

        return newGate

class AdjGate(Gate):
    def __init__(self, gate):
        self.origGate = gate
        self._condense()

    def __str__(self):
        return "("+str(self.origGate)+")'"

    def _condense(self):
        self.mat = self.origGate.mat.H

    def _op(self, obj):
        return Gate([self, obj])

    def shorten(self):
        wrap = Gate([self])
        return wrap.shorten()

class BaseGate(Gate):
    def __init__(self, mat, symb):
        self.childGates = None
        self.mat = mat
        self.symb = symb

    def __str__(self):
        return self.symb

    def _condense(self):
        return False

    def _op(self, obj):
        return Gate([self, obj])

    def shorten(self):
        wrap = Gate([self])
        return wrap.shorten()

class I(BaseGate):
    def __init__(self):
        super().__init__(np.matrix([[1,0],[0,1]]), "I")

class H(BaseGate):
    def __init__(self):
        q = 1/np.sqrt(2)
        super().__init__(np.matrix([[q,q],[q,-1*q]]), "H")

class X(BaseGate):
    def __init__(self):
        super().__init__(np.matrix([[0,1],[1,0]]), "X")

class Y(BaseGate):
    def __init__(self):
        i = complex(0, 1)
        super().__init__(np.matrix([[0,-i],[i,0]]), "Y")

class Z(BaseGate):
    def __init__(self):
        super().__init__(np.matrix([[1,0],[0,-1]]), "Z")

class T(BaseGate):
    def  __init__(self, angle):
        super().__init__(np.matrix([[np.cos(angle),-1*np.sin(angle)],[np.sin(angle),np.cos(angle)]]), "T{"+str(angle)+"}")

class State():
    def __init__(self, array):
        self.vec = array

    def __call__(self,*arg):
        if len(arg)==0:
            return self
        elif len(arg)==1:
            return self._mul(inp)
        else:
            raise Exception("NaBrO")

    def __mul__(self, other):
        return self._mul(other)

    def __repr__(self):
        return "<qState ["+str(self)+"]>"

    def __str__(self):
        return "("+str(self.vec[0])+"|0> + "+str(self.vec[1])+"|1>)"

    def __int__(self):
        return int(0 + 1*(random.random() < self.vec[1]**2))

    def _mul(self, other):
        if isinstance(other, Gate):
            return State(np.dot(self.vec, other.mat))
        elif isinstance(other, State):
            return np.dot(self.vec, other.vec)

    def measure(self):
        if random.random() < self.vec[0]**2:
            return Zero()
        else:
            return One()

class Zero(State):
    def __init__(self):
        super().__init__(np.array([1,0]))

class One(State):
    def __init__(self):
        super().__init__(np.array([0,1]))

I, H, X, Y, Z, One, Zero = I(), H(), X(), Y(), Z(), One(), Zero()