import random

import numpy as np
import pyximport
from numba import jit

pyximport.install()
from ckron import cythkrn

# Raw matrixes to be used for initialization of qubits and gates

# |0> and |1>
__0 = [[1.],
       [0.]]

__1 = [[0.],
       [1.]]

# |+> and |->
__p = [[1 / np.sqrt(2)],
       [1 / np.sqrt(2)]]
__m = [[1 / np.sqrt(2)],
       [-1 / np.sqrt(2)]]

# Gate/Operator raws
_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):
    """Represents a quantum state"""

    def __init__(self, matrix_state):
        """
        Can be initialized with a matrix, e.g. for |0> this is [[0],[1]]
        :param matrix_state: a matrix representing the quantum state
        """
        self.matrix_state = np.array(matrix_state)

    def __getitem__(self, item):
        """
        Kind of hacky way to store a substate of an item for use in Gate
        operations
        so that one can use state[0] and the encompassing operator can access
        this.
        :param item:
        :return:
        """
        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 __eq__(self, other):
        return np.array_equal(self.matrix_state, other.matrix_state)

    def density_matrix(self):
        return self.matrix_state.dot(np.transpose(self.matrix_state))

    def measure_probability(self):
        """
        In a qbit [a, b] normalized: |a|^2 + |b|^2 = 1
        Probability of 0 is |a|^2 and 1 with prob |b|^2

        :returns: tuple of probabilities to measure 0 or 1"""
        return [np.abs(qbit[0]) ** 2 for qbit in self.matrix_state]

    def measure(self):
        """
        This gets a random choice of either 0 and 1 with weights
        based on the probabilities of the qbit

        :returns: classical bit based on qbit probabilities"""
        weights = self.measure_probability()
        format_str = "{:0" + str(int(np.ceil(np.log2(len(weights))))) + "b}"
        choices = [format_str.format(i) for i in range(len(weights))]
        return random.choices(choices, weights)[0]

    def partial_measure(self):
        """
        Say we have the state a|00> + b|01> + c|10> + d|11>.
        Measuring \0> on q1:
          |0>(a|0> + b|1>) + |1>(c|0> + d|1>) ->

        We need to normalize for the other qubit: (a|0> + b|1>) / math.sqrt(|a|^2 + \b|^2)
        :return:
        """
        ...


@jit()
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)
        # TODO: Further optimization with cython - but right now it takes a double only
        rv = cythkrn.kron(rv, item)
    return rv


class Gate(State):
    def on(self, other_state):
        """
        Applies a gate operation on a state.
        If applying to a substate, use the index of the substate, e.g. `H.on(
        bell[0])` will apply the Hadamard gate
        on the 0th qubit of the `bell` state.
        :param other_state: another state (e.g. `H.on(q1)` or a list of
        states (e.g. for `CNOT.on([q1, q2])`)
        :return: the state after the application of the Gate
        """
        this_state_m = self.matrix_state
        if type(other_state) == list:
            other_state = State(kron([e.matrix_state for e in other_state]))
        other_state_m = other_state.matrix_state
        if this_state_m.shape[1] != other_state_m.shape[0]:
            # The two arrays are different sizes
            # Use the Kronicker product of Identity with the state.item where
            # state.item is the substate
            larger_side = max(len(self), len(other_state))
            _list = [this_state_m if i == other_state.item else _I
                     for i in range(larger_side)]
            this_state_m = kron(_list)
        return State(this_state_m.dot(other_state_m))


class Qubit(State): ...


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

CNOT = Gate(_CNOT)


def run_qbit_tests():
    # asserts are sets of tests to check if mathz workz
    # initialize qubits
    _0 = Qubit(__0)
    _1 = Qubit(__1)
    _p = Qubit(__p)
    _m = Qubit(__m)

    # Identity: verify that I|0> == |0> and I|1> == |0>
    assert I.on(_0) == _0
    assert I.on(_1) == _1

    # Pauli X: verify that X|0> == |1> and X|1> == |0>
    assert X.on(_0) == _1
    assert X.on(_1) == _0

    # measure probabilities in sigma_x of |0> and |1>
    # using allclose since dealing with floats
    assert np.allclose(_0.measure_probability(), (1.0, 0.0))
    assert np.allclose(_1.measure_probability(), (0.0, 1.0))

    # applying Hadamard puts the qbit in orthogonal +/- basis
    assert np.array_equal(H.on(_0), _p)
    assert np.array_equal(H.on(_1), _m)

    # measure probabilities in sigma_x of |+> and |->
    # using allclose since dealing with floats
    assert np.allclose(_p.measure_probability(), (0.5, 0.5))
    assert np.allclose(_m.measure_probability(), (0.5, 0.5))


if __name__ == "__main__":
    run_qbit_tests()