import numpy as np

__0 = [[1],
       [0]]

__1 = [[0],
       [1]]

_I = [
    [1, 0],
    [0, 1],
]

_X = [
    [0, 1],
    [1, 0],
]

_Y = [
    [0, complex(0, -1)],
    [complex(0, 1), 0],
]

_Z = [
    [1, 0],
    [0, -1],
]

_H = [
    [1 / np.sqrt(2), 1 / np.sqrt(2)],
    [1 / np.sqrt(2), -1 / np.sqrt(2)]
]

_CNOT = [
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 0, 1],
    [0, 0, 1, 0],
]


class State(object):
    def __init__(self, matrix_state):
        self.matrix_state = np.array(matrix_state)

    def __getitem__(self, item):
        if item >= len(self):
            raise IndexError
        self.item = item
        return self

    def __len__(self):
        return int(np.log2(self.matrix_state.shape[0]))

    def __repr__(self):
        return str(self.matrix_state)


def kron(_list):
    """Calculates a Kronicker product of a list of matrices
    This is essentially a np.kron(*args) since np.kron takes (a,b)"""
    if type(_list) != list or len(_list) <= 1:
        return np.array([])
    rv = np.array(_list[0])
    for item in _list[1:]:
        rv = np.kron(rv, item)
    return rv


class Gate(State):
    def on(self, state):
        """
        Applies
        :param state: another state (e.g. H(q1) or a list of states (e.g. for CNOT([q1, q2]))
        :return:
        """
        this_state = self.matrix_state
        if type(state) == list:
            other_state = kron([e.matrix_state for e in state])
        else:
            other_state = state.matrix_state
        if this_state.shape[1] != other_state.shape[0]:
            # The two arrays are different sizes
            # Use the Kronicker product of Identity with the state.item
            larger_side = max(this_state.shape[1], this_state.shape[0])
            _list = [this_state if i == state.item else _I for i in range(larger_side)]
            this_state = kron(_list)
        return State(this_state.dot(other_state))


class Qubit(State): ...


# List of gates
I = Gate(_I)
X = Gate(_X)
Y = Gate(_Y)
Z = Gate(_Z)
H = Gate(_H)

CNOT = Gate(_CNOT)