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, times=10):
    """
    Measure result from computation based on the qubit population
    """
    probabilites = [q[0] ** 2 for q in qubits]
    return [bin(m) for m in random.choices(
        population=range(len(qubits)),
        weights=probabilites,
        k=times
    )]


_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
    # Put two qubits in superposition
    superposition = Qubit(H * _0) * Qubit(H * _0)
    print(superposition)
    # Entangle the qubits
    entanglement = Qubit(CNOT * superposition)
    print(entanglement)

    # TODO: What does it mean to apply H gate on the entangled qubits???
    #  https://en.wikipedia.org/wiki/Controlled_NOT_gate#Behaviour_in_the_Hadamard_transformed_basis
    #  Should it be:
    #    q1 = Qubit(H * entanglement[:2])
    #    q2 = Qubit(H * entanglement[2:])
    #  But how does one get the 4x4 matrix of CNOT???
    #
    # q1 = Qubit(H * Qubit(entanglement[:2]))
    # print(q1)
    # q2 = Qubit(H * Qubit(entanglement[2:]))
    # print(q2)
    # print(q1 * q2)

    # this should get us 1/4 chance for each combination of 00 01 10 and 10;
    # measure 42 times
    print(measure(entanglement, times=42))

    # 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()