import random
from math import sqrt, log


class Gate(list):
    # Matrix multiplication
    def __mul__(self, other):
        rv = Gate()
        for kol in range(len(other[0])):
            for _ in range(len(self)):
                rv.append([0])
        # iterate through rows of a
        for i in range(len(self)):
            # iterate through columns of b
            for j in range(len(other[0])):
                # iterate through rows of b
                for k in range(len(other)):
                    rv[i][j] += self[i][k] * other[k][j]
        return rv


class Qubit(list):
    # tensor multiplication
    def __mul__(self, other):
        rv = Qubit()
        for i in self:
            for j in other:
                rv.append([i[0] * j[0]])
        return rv


def measure(qubits):
    """
    Measure result from computation
    """
    probabilites = [q[0] ** 2 for q in qubits]
    return bin(random.choices(
        population=range(len(qubits)),
        weights=probabilites
    )[0])


_0 = Qubit([[1], [0]])
_1 = Qubit([[0], [1]])

I = Gate([
    [1, 0],
    [0, 1],
])
X = Gate([
    [0, 1],
    [1, 0],
])
Y = Gate([
    [0, complex(0, -1)],
    [complex(0, 1), 0],
])
Z = Gate([
    [1, 0],
    [0, -1],
])
CNOT = Gate([
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 0, 1],
    [0, 0, 1, 0],
])
H = Gate([
    [1 / sqrt(2), 1 / sqrt(2)],
    [1 / sqrt(2), -1 / sqrt(2)],
])


def main():
    # TEST IDENTITY
    assert _1 == I * _1
    assert _0 == I * _0

    # TEST BIT-FLIP
    assert _1 == X * _0
    assert _0 == X * _1

    # TEST CNOT GATE -  Control * Target
    # Control is FALSE, so target doesn't change
    assert CNOT * (_0 * _0) == (_0 * _0)
    assert CNOT * (_0 * _1) == (_0 * _1)
    # Control is TRUE, so target MUST change
    assert CNOT * (_1 * _0) == (_1 * _1)
    assert CNOT * (_1 * _1) == (_1 * _0)

    # TEST HADAMARD
    superposition = Qubit(H * _0) * Qubit(H * _0)
    print(superposition)
    CNOT_superposition = Qubit(CNOT * superposition)
    print(CNOT_superposition)

    print(measure(CNOT_superposition))
    # print(INVERTED_CNOT)

    # TEST MY MEMORY
    # rv = _0
    # QUBITS = 23
    # print(2**QUBITS)
    # print("{:.1f} KB".format(2**QUBITS / 2**10))
    # print("{:.1f} MB".format(2**QUBITS / 2**20))
    # print("{:.1f} GB".format(2**QUBITS / 2**30))
    # for i in range(QUBITS):
    #     rv *= _0
    # print(len(rv))


if __name__ == "__main__":
    main()