import numpy as np


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

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


class Gate(State):
    def on(self, state, q2=None):
        this = self._matrix_state
        a = state._matrix_state
        if q2:
            a = np.kron(a, q2._matrix_state)
        return State(this.dot(a))


class Qubit(State): ...


_0 = np.array([[1],
               [0]])

H = Gate(np.array([
    [1 / np.sqrt(2), 1 / np.sqrt(2)],
    [1 / np.sqrt(2), -1 / np.sqrt(2)]
]))

CNOT = Gate(np.array([
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 0, 1],
    [0, 0, 1, 0],
]))

if __name__ == "__main__":
    # E.g. Bell state
    q1 = Qubit(_0)
    q2 = Qubit(_0)
    bell = CNOT.on(H.on(q1), q2)
    print(bell)