import random
from collections import defaultdict
from functools import reduce
from numbers import Complex
from typing import Union

import numpy as np
import sympy as sp

from load_test import sizeof_fmt

ListOrNdarray = Union[list, np.ndarray]

REPR_TENSOR_OP = "⊗"
REPR_GREEK_BETA = "β"
REPR_GREEK_PSI = "ψ"
REPR_GREEK_PHI = "φ"
REPR_MATH_KET = "⟩"
REPR_MATH_BRA = "⟨"
REPR_MATH_SQRT = "√"
REPR_MATH_SUBSCRIPT_NUMBERS = "₀₁₂₃₄₅₆₇₈₉"

# Keep a reference to already used states for naming
UNIVERSE_STATES = []


class Matrix(object):
    """Wraps a Matrix... it's for my understanding, this could easily probably be np.array"""

    def __init__(self, m: ListOrNdarray = None, *args, **kwargs):
        """
        Can be initialized with a matrix, e.g. for |0> this is [[0],[1]]
        :param m: a matrix representing the quantum state
        :param name: the name of the state (optional)
        """
        if m is None:
            self.m = np.array([])
        elif type(m) is list:
            self.m = np.array(m)
        elif type(m) is np.ndarray:
            self.m = m
        else:
            raise TypeError("m needs to be a list or ndarray type")

    def __add__(self, other):
        return Matrix(self.m + other.m)

    def __eq__(self, other):
        if isinstance(other, Complex):
            return bool(np.allclose(self.m, other))
        elif isinstance(other, TwoQubitPartial):
            return False
        return np.allclose(self.m, other.m)

    def __mul_or_kron__(self, other):
        """
        Multiplication with a number is Linear op;
        with another state is a composition via the Kronecker/Tensor product
        """
        if isinstance(other, Complex):
            return Matrix(self.m * other)
        elif isinstance(other, Matrix):
            return self.__class__(np.kron(self.m, other.m))
        raise NotImplementedError

    def __rmul__(self, other):
        return self.__mul_or_kron__(other)

    def __mul__(self, other):
        return self.__mul_or_kron__(other)

    def __or__(self, other):
        """Define inner product: <self|other>
        """
        m = np.dot(self._conjugate_transpose(), other.m)
        try:
            return self.__class__(m)
        except:
            try:
                return other.__class__(m)
            except:
                ...
        return Matrix(m)

    def __len__(self):
        return len(self.m)

    def outer(self, other):
        """Define outer product |0><0|"""
        return Matrix(np.outer(self.m, other.m))

    def x(self, other):
        """Define outer product |0><0| looks like |0x0| which is 0.x(0)"""
        return self.outer(other)

    def _conjugate_transpose(self):
        return self.m.transpose().conjugate()

    def _complex_conjugate(self):
        return self.m.conjugate()

    def conjugate_transpose(self):
        return Matrix(self._conjugate_transpose())

    def complex_conjugate(self):
        return Matrix(self._complex_conjugate())

    @property
    def human_m(self):
        return humanize(self.m)

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


class Vector(Matrix):
    def __init__(self, m: ListOrNdarray = None, *args, **kwargs):
        super().__init__(m, *args, **kwargs)
        if not self._is_vector():
            raise TypeError("Not a vector")

    def _is_vector(self):
        return self.m.shape[1] == 1


class State(Vector):
    def __init__(self, m: ListOrNdarray = None, name: str = '', *args, **kwargs):
        """State vector representing quantum state"""
        super().__init__(m, *args, **kwargs)
        self.name = name
        if not self._is_normalized():
            raise TypeError("Not a normalized state vector")

    def _is_normalized(self):
        return np.isclose(np.sum(np.abs(self.m ** 2)), 1.0)

    @staticmethod
    def _normalize_angles(theta, phi):
        # theta is between [0 and pi]
        theta = theta.real % np.pi if theta.real != np.pi else theta.real
        # phi is between (0 and 2*pi]
        phi = phi.real % (2 * np.pi)
        return theta, phi

    @classmethod
    def from_angles(cls, theta, phi):
        theta, phi = cls._normalize_angles(theta, phi)
        m0 = np.cos(theta / 2)
        m1 = np.sin(theta / 2) * np.power(np.e, (1j * phi))
        m = m0 * _0 + m1 * _1
        return cls(m.m)

    def to_angles(self):
        if not self.m.shape == (2, 1):
            raise Exception("State needs to be 2x1 matrix")
        m0, m1 = self.m[0][0], self.m[1][0]

        # theta is between 0 and pi
        theta = 2 * np.arccos(m0)
        if theta < 0:
            theta += np.pi
        assert 0 <= theta <= np.pi

        div = np.sin(theta / 2)
        if div == 0:
            # here is doesn't matter what phi is as phase at the poles is arbitrary
            phi = 0
        else:
            exp = m1 / div
            phi = np.log(complex(exp)) / 1j
            if phi < 0:
                phi += 2 * np.pi
        assert 0 <= phi <= 2 * np.pi
        return theta, phi

    def rotate_x(self, theta):
        return Rx(theta).on(self)

    def rotate_y(self, theta):
        return Ry(theta).on(self)

    def rotate_z(self, theta):
        return Rz(theta).on(self)

    def __repr__(self):
        if not self.name:
            if np.count_nonzero(self.m == 1):
                fmt_str = self.get_fmt_of_element()
                index = np.where(self.m == [1])[0][0]
                self.name = fmt_str.format(index)
            else:
                for well_known_state in well_known_states:
                    try:
                        if self == well_known_state:
                            return repr(well_known_state)
                    except:
                        ...
                for used_state in UNIVERSE_STATES:
                    try:
                        if self == used_state:
                            return repr(used_state)
                    except:
                        ...
                next_state_sub = ''.join([REPR_MATH_SUBSCRIPT_NUMBERS[int(d)] for d in str(len(UNIVERSE_STATES))])
                self.name = '{}{}'.format(REPR_GREEK_PSI, next_state_sub)
        UNIVERSE_STATES.append(self)
        # matrix_rep = "{}".format(self.m).replace('[', '').replace(']', '').replace('\n', '|').strip()
        # state_name = '|{}> = {}'.format(self.name, matrix_rep)
        state_name = '|{}>'.format(self.name)
        return state_name

    def norm(self):
        """Norm/Length of the vector = sqrt(<self|self>)"""
        return self.length()

    def length(self):
        """Norm/Length of the vector = sqrt(<self|self>)"""
        return np.sqrt((self | self).m).item(0)

    def is_orthogonal(self, other):
        """If the inner (dot) product is zero, this vector is orthogonal to other"""
        return self | other == 0

    def get_prob(self, j):
        """pr(j) = |<e_j|self>|^2"""
        # j_th basis vector
        e_j = State([[1] if i == int(j) else [0] for i in range(len(self.m))])
        return np.absolute((e_j | self).m.item(0)) ** 2

    def get_fmt_of_element(self):
        return "{:0" + str(int(np.ceil(np.log2(len(self))))) + "b}"

    def measure(self):
        """
        Measures in the computational basis
        TODO: Generalize the method so it takes a basis
        TODO: Should we memoize the result? Should we memoize per basis?
              If it's measured twice, should it return the same state?
              What about if measured twice but in different bases?
              E.g. measure1 -> computation -> A
                   measure2 -> +/- basis -> B
                   measure3 -> computation -> should it return A again or random weighted?
                   measure4 -> +/- basis -> should it return B again or random weighted?
        :return: binary representation of the measured qubit (e.g. "011")
        """
        weights = [self.get_prob(j) for j in range(len(self))]
        format_str = self.get_fmt_of_element()
        choices = [format_str.format(i) for i in range(len(weights))]
        return random.choices(choices, weights)[0]

    def measure_with_op(self, mo):
        """
        Measures with a measurement operator mo
        TODO: Can't define `mo: MeasurementOperator` because in python you can't declare classes before defining them
        """
        m = mo.on(self) / np.sqrt(mo.get_prob(self))
        return State(m)

    def measure_n(self, n=1):
        """measures n times
        TODO: Does this make sense? When a state is measured once, it should "collapse"
        """
        measurements = defaultdict(int)
        for _ in range(n):
            k = self.measure()
            measurements[k] += 1
        return measurements

    def pretty_print(self):
        format_str = self.get_fmt_of_element() + " | {}"

        for i, element in enumerate(self.m):
            print(format_str.format(i, humanize_num(element[0])))

    @classmethod
    def from_string(cls, q):
        try:
            bin_repr = q[1:-1]
            dec_bin = int(bin_repr, 2)
            bin_length = 2 ** len(bin_repr)
            arr = [[1] if i == dec_bin else [0] for i in range(bin_length)]
        except:
            raise Exception("State from string should be of the form |00..01> with all numbers either 0 or 1")
        return cls(arr)

    @staticmethod
    def get_random_state_angles():
        phi = np.random.uniform(0, 2 * np.pi)
        t = np.random.uniform(0, 1)
        theta = np.arccos(1 - 2 * t)
        return theta, phi

    def get_bloch_coordinates(self):
        theta, phi = self.to_angles()
        x = np.sin(theta) * np.cos(phi)
        y = np.sin(theta) * np.sin(phi)
        z = np.cos(theta)
        return [x, y, z]


def s(q):
    """Helper method for creating state easily"""
    if type(q) == str:
        # e.g. |000>
        if q[0] == '|' and q[-1] == '>':
            return State.from_string(q)
    elif type(q) == list:
        # e.g. s([1,0]) => |0>
        return State(np.reshape(q, (len(q), 1)))
    return State(q)


def humanize_num(fl, tolerance=1e-3):
    if np.abs(fl.imag) < tolerance:
        fl = fl.real
    if np.abs(fl.real) < tolerance:
        fl = 0 + 1j * fl.imag
    try:
        return sp.nsimplify(fl, [sp.pi], tolerance, full=True)
    except:
        return sp.nsimplify(fl, [sp.pi], tolerance)


def pp(m, tolerance=1e-3):
    for element in m:
        print(humanize_num(element, tolerance=tolerance))


def humanize(m):
    rv = []
    for element in m:
        if type(element) in [np.ndarray, list]:
            rv.append(humanize(element))
        else:
            rv.append(humanize_num(element))
    return rv


class Operator(object):
    def __init__(self, func=None, *args, **kwargs):
        """An Operator turns one function into another"""
        self.func = func

    def on(self, *args, **kwargs):
        return self(*args, **kwargs)

    def __call__(self, *args, **kwargs):
        return self.func(*args, **kwargs)


class LinearOperator(Operator):
    def __init__(self, func=None, *args, **kwargs):
        """Linear operators satisfy f(x+y) = f(x) + f(y) and a*f(x) = f(a*x)"""
        super().__init__(func, *args, **kwargs)
        if not self._is_linear():
            raise TypeError("Not a linear operator")

    def _is_linear(self):
        # TODO: How to verify if the func is linear?
        #       in case of Unitary Operator, self.func is a lambda that takes a Matrix (assumes has .m component)
        return True
        # a, b = sympy.symbols('a, b')
        # expr, vars_ = a+b, [a, b]
        # for x in vars_:
        #     for y in vars_:
        #         try:
        #             if not sympy.Eq(sympy.diff(expr, x, y), 0):
        #                 return False
        #         except TypeError:
        #             return False
        # return True


class SquareMatrix(Matrix):
    def __init__(self, m: ListOrNdarray, *args, **kwargs):
        super().__init__(m, *args, **kwargs)
        if not self._is_square():
            raise TypeError("Not a Square matrix")

    def _is_square(self):
        return self.m.shape[0] == self.m.shape[1]


class UnitaryMatrix(SquareMatrix):
    def __init__(self, m: ListOrNdarray, *args, **kwargs):
        """Represents a Unitary matrix that satisfies UU+ = I"""
        super().__init__(m, *args, **kwargs)
        if not self._is_unitary():
            raise TypeError("Not a Unitary matrix")

    def _is_unitary(self):
        """Checks if the matrix product of itself with conjugate transpose is Identity UU+ = I"""
        UU_ = np.dot(self._conjugate_transpose(), self.m)
        I = np.eye(self.m.shape[0])
        return np.isclose(UU_, I).all()


class HermitianMatrix(SquareMatrix):
    def __init__(self, m: ListOrNdarray, *args, **kwargs):
        """Represents a Hermitian matrix that satisfies U=U+"""
        super().__init__(m, *args, **kwargs)
        if not self._is_hermitian():
            raise TypeError("Not a Hermitian matrix")

    def _is_hermitian(self):
        """Checks if its equal to the conjugate transpose U = U+"""
        return np.isclose(self.m, self._conjugate_transpose()).all()


class UnitaryOperator(LinearOperator, UnitaryMatrix):
    def __init__(self, m: ListOrNdarray, name: str = '', *args, **kwargs):
        """UnitaryOperator inherits from both LinearOperator and a Unitary matrix
        It is used to act on a State vector by defining the operator to be the dot product"""
        self.name = name
        UnitaryMatrix.__init__(self, m=m, *args, **kwargs)
        LinearOperator.__init__(self, func=self.operator_func, *args, **kwargs)

    def operator_func(self, other):
        return State(np.dot(self.m, other.m))

    def __repr__(self):
        if self.name:
            return '-{}-'.format(self.name)
        return str(self.m)


class HermitianOperator(LinearOperator, HermitianMatrix):
    def __init__(self, m: ListOrNdarray, name: str = '', *args, **kwargs):
        """HermitianMatrix inherits from both LinearOperator and a Hermitian matrix
        It is used to act on a State vector by defining the operator to be the dot product"""
        self.name = name
        HermitianMatrix.__init__(self, m=m, *args, **kwargs)
        LinearOperator.__init__(self, func=self.operator_func, *args, **kwargs)

    def operator_func(self, other):
        """This might not return a normalized state vector so don't wrap it in State"""
        return np.dot(self.m, other.m)

    def __repr__(self):
        if self.name:
            return '-{}-'.format(self.name)
        return str(self.m)


# How to add a CNOT gate to the Quantum Processor?
# Imagine if I have to act on a 3-qubit computer and CNOT(q1, q3)
#
# Decomposed CNOT :
# reverse engineered from
# https://quantumcomputing.stackexchange.com/questions/4252/how-to-derive-the-cnot-matrix-for-a-3-qbit-system-where-the-control-target-qbi
#
# CNOT(q1, I, q2):
# |0><0| x I_2 x I_2 + |1><1| x I_2 x X
# np.kron(np.kron(np.outer(_0.m, _0.m), np.eye(2)), np.eye(2)) + np.kron(np.kron(np.outer(_1.m, _1.m), np.eye(2)), X.m)
#
#
# CNOT(q1, q2):
# |0><0| x I + |1><1| x X
# np.kron(np.outer(_0.m, _0.m), np.eye(2)) + np.kron(np.outer(_1.m, _1.m), X.m)
# _0.x(_0) * Matrix(I.m) + _1.x(_1) * Matrix(X.m)

class TwoQubitPartial(object):
    def __init__(self, rpr):
        self.rpr = rpr
        self.operator = None

    def __repr__(self):
        return str("-{}-".format(self.rpr))


C_partial = TwoQubitPartial("C")
x_partial = TwoQubitPartial("x")


class TwoQubitOperator(UnitaryOperator):
    def __init__(self, m: ListOrNdarray, A: TwoQubitPartial, B: TwoQubitPartial,
                 A_p: UnitaryOperator, B_p: UnitaryOperator, *args, **kwargs):
        super().__init__(m, *args, **kwargs)
        A.operator, B.operator = self, self
        self.A = A
        self.B = B
        self.A_p = A_p
        self.B_p = B_p

        A.operator = self
        B.operator = self

    def verify_step(self, step):
        if not (step.count(self.A) == 1 and step.count(self.B) == 1):
            raise RuntimeError("Both {} and {} need to be defined in the same step exactly once".format(
                self.A, self.B
            ))

    def compose(self, step, state):
        # _0.x(_0) * Matrix(I.m) + _1.x(_1) * Matrix(X.m)
        outer_0, outer_1 = [], []
        for s in step:
            if s == self.A:
                outer_0.append(_0.x(_0))
                outer_1.append(_1.x(_1))
            elif s == self.B:
                outer_0.append(Matrix(self.A_p.m))
                outer_1.append(Matrix(self.B_p.m))
            else:
                outer_0.append(Matrix(s.m))
                outer_1.append(Matrix(s.m))
        return reduce((lambda x, y: x * y), outer_0) + reduce((lambda x, y: x * y), outer_1)


"""
Define States and Operators 
"""

_0 = State([[1],
            [0]],
           name='0')

_1 = State([[0],
            [1]],
           name='1')

# |+> and |-> states
_p = State([[1 / np.sqrt(2)],
            [1 / np.sqrt(2)]],
           name='+')

_m = State([[1 / np.sqrt(2)],
            [- 1 / np.sqrt(2)]],
           name='-')

# 4 Bell states
b_00 = State([[1 / np.sqrt(2)],
              [0],
              [0],
              [1 / np.sqrt(2)]],
             name='B00')

b_01 = State([[0],
              [1 / np.sqrt(2)],
              [1 / np.sqrt(2)],
              [0]],
             name='B01')

b_10 = State([[1 / np.sqrt(2)],
              [0],
              [0],
              [-1 / np.sqrt(2)]],
             name='B10')

b_11 = State([[0],
              [1 / np.sqrt(2)],
              [-1 / np.sqrt(2)],
              [0]],
             name='B11')

well_known_states = [_p, _m, b_00, b_01, b_10, b_11]

_ = I = UnitaryOperator([[1, 0],
                         [0, 1]],
                        name="-")

X = UnitaryOperator([[0, 1],
                     [1, 0]],
                    name="X")

Y = UnitaryOperator([[0, -1j],
                     [1j, 0]],
                    name="Y")

Z = UnitaryOperator([[1, 0],
                     [0, -1]],
                    name="Z")


# These are rotations that are specified commonly e.g. in
#       https://www.quantum-inspire.com/kbase/rotation-operators/
#       http://www.vcpc.univie.ac.at/~ian/hotlist/qc/talks/bloch-sphere-rotations.pdf
#       and elsewhere DO NOT equate X, Y and Z for theta=np.pi
# However - they are correct up to a global phase which is all that matters for measurement purposes
#

def Rx(theta):
    return UnitaryOperator([[np.cos(theta / 2), -1j * np.sin(theta / 2)],
                            [-1j * np.sin(theta / 2), np.cos(theta / 2)]],
                           name="Rx")


def Ry(theta):
    return UnitaryOperator([[np.cos(theta / 2), -np.sin(theta / 2)],
                            [np.sin(theta / 2), np.cos(theta / 2)]],
                           name="Ry")


def Rz(theta):
    return UnitaryOperator([[np.power(np.e, -1j * theta / 2), 0],
                            [0, np.power(np.e, 1j * theta / 2)]],
                           name="Rz")

H = UnitaryOperator([[1 / np.sqrt(2), 1 / np.sqrt(2)],
                     [1 / np.sqrt(2), -1 / np.sqrt(2)], ],
                    name="H")

CNOT = TwoQubitOperator([[1, 0, 0, 0],
                         [0, 1, 0, 0],
                         [0, 0, 0, 1],
                         [0, 0, 1, 0], ],
                         C_partial, x_partial, I, X)


def assert_raises(exception, msg, callable, *args, **kwargs):
    try:
        callable(*args, **kwargs)
    except exception as e:
        assert str(e) == msg
        return
    assert False


def assert_not_raises(exception, msg, callable, *args, **kwargs):
    try:
        callable(*args, **kwargs)
    except exception as e:
        if str(e) == msg:
            assert False


def test_unitary_hermitian():
    # Unitary is UU+ = I; Hermitian is U = U+
    # Matrixes could be either, neither or both
    # Quantum operators (gates) are described *only* by unitary transformations
    # Hermitian operators are used for measurement operators - https://towardsdatascience.com/understanding-basics-of-measurements-in-quantum-computation-4c885879eba0
    h_not_u = [
        [1, 0],
        [0, 2],
    ]
    assert_not_raises(TypeError, "Not a Hermitian matrix", HermitianMatrix, h_not_u)
    assert_raises(TypeError, "Not a Unitary matrix", UnitaryMatrix, h_not_u)

    u_not_h = [
        [1, 0],
        [0, 1j],
    ]
    assert_raises(TypeError, "Not a Hermitian matrix", HermitianMatrix, u_not_h)
    assert_not_raises(TypeError, "Not a Unitary matrix", UnitaryMatrix, u_not_h)

    u_and_h = [
        [0, 1],
        [1, 0],
    ]
    assert_not_raises(TypeError, "Not a Hermitian matrix", HermitianMatrix, u_and_h)
    assert_not_raises(TypeError, "Not a Unitary matrix", UnitaryMatrix, u_and_h)

    not_u_not_h = [
        [1, 2],
        [0, 1],
    ]
    assert_raises(TypeError, "Not a Hermitian matrix", HermitianMatrix, not_u_not_h)
    assert_raises(TypeError, "Not a Unitary matrix", UnitaryMatrix, not_u_not_h)


class MeasurementOpeartor(HermitianOperator):
    """Measurement operators are Hermitians: <ψ|M†_m M_m|ψ>"""

    @classmethod
    def create_from_prob(cls, matrix: Matrix):
        """returns |M†_m><M_m|"""
        return cls(matrix.conjugate_transpose().x(matrix).m)

    def get_prob(self, state: State):
        state_ct = state.conjugate_transpose()
        return state_ct.m.dot(self.m.dot(state.m)).item()


def test_measurement_ops():
    m0 = MeasurementOpeartor.create_from_prob(Matrix([1, 0]))
    m1 = MeasurementOpeartor.create_from_prob(Matrix([0, 1]))
    assert m0 == Matrix([[1, 0],
                         [0, 0]])
    assert m1 == Matrix([[0, 0],
                         [0, 1]])

    # p(0) -> probability of measurement to yield a 0
    assert m0.get_prob(_0) == 1.0
    assert m1.get_prob(_0) == 0.0
    assert m0.get_prob(_1) == 0.0
    assert m1.get_prob(_1) == 1.0

    # Post-state measurement of qubit with operator
    assert _p.measure_with_op(m0) == _0
    assert _p.measure_with_op(m1) == _1
    assert _m.measure_with_op(m0) == _0
    assert _m.measure_with_op(m1) == s([0, -1])


def abs_squared(x):
    return np.abs(x) ** 2


def testRotPauli():
    # From http://www.vcpc.univie.ac.at/~ian/hotlist/qc/talks/bloch-sphere.pdf
    # / slide 11:
    # However, the only measurable quantities are the probabilities
    # |α|**2 and |β|**2, so multiplying the state by an arbitrary factor
    # exp(iγ) (a global phase) has no observable consequences
    assert np.allclose(abs_squared(Rx(np.pi).m), abs_squared(X.m))
    assert np.allclose(abs_squared(Ry(np.pi).m), abs_squared(Y.m))
    assert np.allclose(abs_squared(Rz(np.pi).m), abs_squared(Z.m))


def test():
    # Test properties of Hilbert vector space
    # The four postulates of Quantum Mechanics
    # I: States | Associated to any physical system is a complex vector space
    #             known as the state space of the system. If the system is closed
    #             then the system is described completely by its state vector
    #             which is a unit vector in the space.
    # Mathematically, this vector space is also a function space
    assert _0 + _1 == _1 + _0  # commutativity of vector addition
    assert _0 + (_1 + _p) == (_0 + _1) + _p  # associativity of vector addition
    assert 8 * (_0 + _1) == 8 * _0 + 8 * _1  # Linear when multiplying by constants
    assert _0 | _0 == 1  # parallel have 1 product
    assert _0 | _1 == 0  # orthogonal have 0 product
    assert _0.is_orthogonal(_1)
    assert _1 | (8 * _0) == 8 * (_1 | _0)  # Inner product is linear multiplied by constants
    assert _p | (_1 + _0) == (_p | _1) + (_p | _0)  # Inner product is linear in superpos of vectors

    assert np.isclose(_1.length(), 1.0)  # all of the vector lengths are normalized
    assert np.isclose(_0.length(), 1.0)
    assert np.isclose(_p.length(), 1.0)

    assert _0 | _1 == (_1 | _0).complex_conjugate()  # non-commutative inner product

    test_to_from_angles()

    # II: Dynamics | The evolution of a closed system is described by a unitary transformation
    #
    # Operators turn one vector into another
    # the times 2 operator should return the times two multiplication
    _times_2 = Operator(lambda x: 2 * x)
    assert _times_2.on(5) == 10
    assert _times_2(5) == 10

    # Understanding the difference between unitary and hermitians
    test_unitary_hermitian()

    # Pauli X gate flips the |0> to |1> and the |1> to |0>
    assert X | _1 == _0
    assert X | _0 == _1

    # Test the Y Pauli operator with complex number literal notation
    assert Y | _0 == State([[0],
                            [1j]])
    assert Y | _1 == State([[-1j],
                            [0]])

    # Test Pauli rotation gates
    testRotPauli()

    # III: Measurement | A quantum measurement is described by an orthonormal basis |e_j>
    #                    for state space. If the initial state of the system is |ψ>
    #                    then we get outcome j with probability pr(j) = |<e_j|ψ>|^2
    # Note: The postulates are applicable on closed, isolated systems.
    #       Systems that are closed and are described by unitary time evolution by a Hamiltonian
    #       can be measured by projective measurements. Systems are not closed in reality and hence
    #       are immeasurable using projective measurements.
    #       POVM (Positive Operator-Valued Measure) is a restriction on the projective measurements,
    #       such that it encompasses everything except the environment.

    assert _0.get_prob(0) == 1  # Probability for |0> in 0 is 1
    assert _0.get_prob(1) == 0  # Probability for |0> in 1 is 0

    assert _1.get_prob(0) == 0  # Probability for |1> in 0 is 0
    assert _1.get_prob(1) == 1  # Probability for |1> in 1 is 1

    assert np.isclose(_p.get_prob(0), 0.5)  # Probability for |+> in 0 is 0.5
    assert np.isclose(_p.get_prob(1), 0.5)  # Probability for |+> in 1 is 0.5

    # Test measurement operators
    test_measurement_ops()

    # IV: Compositing | The state space of a composite physical system
    #                   is the tensor product of the state spaces
    #                   of the component physical systems.
    #                   i.e. if we have systems numbered 1 through n, ψ_i,
    #                   then the joint state is |ψ_1> ⊗ |ψ_2> ⊗ ... ⊗ |ψ_n>
    assert _0 * _0 == State([[1], [0], [0], [0]])
    assert _0 * _1 == State([[0], [1], [0], [0]])
    assert _1 * _0 == State([[0], [0], [1], [0]])
    assert _1 * _1 == State([[0], [0], [0], [1]])

    # CNOT applies a control qubit on a target.
    # If the control is a |0>, target remains unchanged.
    # If the control is a |1>, target is flipped.
    assert CNOT.on(_0 * _0) == _0 * _0
    assert CNOT.on(_0 * _1) == _0 * _1
    assert CNOT.on(_1 * _0) == _1 * _1
    assert CNOT.on(_1 * _1) == _1 * _0

    # ALL FOUR NOW - Create a Bell state (I) that has H|0> (II), measures (III) and composition in CNOT (IV)
    # Bell state
    # First - create a superposition
    H = UnitaryOperator([[1 / np.sqrt(2), 1 / np.sqrt(2)],
                         [1 / np.sqrt(2), -1 / np.sqrt(2)], ])
    superpos = H | _0
    assert superpos == _p

    # Then CNOT the superposition with a |0> qubit
    bell = CNOT | (superpos * _0)
    assert bell == State([[1 / np.sqrt(2)],
                          [0.],
                          [0.],
                          [1 / np.sqrt(2)], ])

    assert np.isclose(bell.get_prob(0b00), 0.5)  # Probability for bell in 00 is 0.5
    assert np.isclose(bell.get_prob(0b01), 0)  # Probability for bell in 01 is 0.
    assert np.isclose(bell.get_prob(0b10), 0)  # Probability for bell in 10 is 0.
    assert np.isclose(bell.get_prob(0b11), 0.5)  # Probability for bell in 11 is 0.5

    ################################
    # TODO: Don't know where outer product fits - something about density operator?
    assert _0.x(_0) == Matrix([[1, 0],
                               [0, 0]])
    assert _0.x(_1) == Matrix([[0, 1],
                               [0, 0]])
    assert _1.x(_0) == Matrix([[0, 0],
                               [1, 0]])
    assert _1.x(_1) == Matrix([[0, 0],
                               [0, 1]])


def test_to_from_angles():
    for q in [_0, _1, _p, _m]:
        angles = q.to_angles()
        s = State.from_angles(*angles)
        assert q == s

    assert State.from_angles(0, 0) == _0
    assert State.from_angles(np.pi, 0) == _1
    assert State.from_angles(np.pi / 2, 0) == _p
    assert State.from_angles(np.pi / 2, np.pi) == _m
    for theta, phi, qbit in [
        (0, 0, _0),
        (np.pi, 0, _1),
        (np.pi / 2, 0, _p),
        (np.pi / 2, np.pi, _m),
    ]:
        s = State.from_angles(theta, phi)
        assert s == qbit


def test_measure_n():
    qq = State([
        [0.5],
        [0.5],
        [0.5],
        [0.5],
    ])
    qq.measure_n(100)


def naive_load_test(N):
    import os
    import psutil
    import gc
    from time import time
    from sys import getsizeof

    print("{:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10}".format(
        "qbits",
        "kron_len",
        "mem_used",
        "mem_per_q",
        "getsizeof",
        "getsiz/len",
        "nbytes",
        "nbytes/len",
        "time"))

    _0 = State([[1], [0]], name='0')

    process = psutil.Process(os.getpid())
    mem_init = process.memory_info().rss
    for i in range(2, N + 1):
        start = time()
        m = _0
        for _ in range(i):
            m = m * _0
        len_m = len(m)
        elapsed = time() - start
        mem_b = process.memory_info().rss - mem_init
        print("{:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10}".format(
            i,
            len(m),
            sizeof_fmt(mem_b),
            sizeof_fmt(mem_b / len_m),
            sizeof_fmt(getsizeof(m)),
            sizeof_fmt(getsizeof(m) / len_m),
            sizeof_fmt(m.m.nbytes),
            sizeof_fmt(m.m.nbytes / len_m),
            np.round(elapsed, 2)))

        gc.collect()


class QuantumCircuit(object):
    def __init__(self, n_qubits: int, initial_steps=None):
        self.n_qubits = n_qubits
        if not initial_steps:
            self.steps = [[_0 for _ in range(n_qubits)], ]
        else:
            self.steps = initial_steps
        self._called_add_row = 0

    @property
    def rows(self):
        return self._called_add_row

    def add_row_step(self, row: int, step: int, qbit_state):
        if len(self.steps) <= step:
            self.steps += [[I for _ in range(self.n_qubits)] for _ in range(len(self.steps) - step + 1)]
        self.steps[step][row] = qbit_state

    def add_step(self, step_data: list):
        if len(step_data) != self.n_qubits:
            raise RuntimeError("Length of step is: {}, should be: {}".format(len(step_data), self.n_qubits))
        step_i = len(step_data)
        for row, qubit_state in enumerate(step_data):
            self.add_row_step(row, step_i, qubit_state)

    def add_steps(self, steps_data: list):
        for step_data in steps_data:
            self.add_step(step_data)

    def add_row(self, row_data: list):
        if self._called_add_row >= self.n_qubits:
            raise RuntimeError("Adding more rows than qubits")
        for step_i, qubit_state in enumerate(row_data):
            self.add_row_step(self._called_add_row, step_i + 1, qubit_state)
        self._called_add_row += 1

    def add_rows(self, rows_data: list):
        for row_data in rows_data:
            self.add_row(row_data)

    def compose_quantum_state(self, step):
        partials = [op for op in step if isinstance(op, TwoQubitPartial)]
        # TODO: No more than 1 TwoQubitGate **OR** UnitaryOperator can be used in a step
        for partial in partials:
            two_qubit_op = partial.operator
            two_qubit_op.verify_step()
            return two_qubit_op.compose(step)
        return reduce((lambda x, y: x * y), step)

    def step(self):
        if self.current_step == 0 and self.current_state == self.HALT_STATE:
            self.current_state = self.RUNNING_STATE
        if self.current_step >= len(self.steps):
            self.current_state = self.HALT_STATE
            raise RuntimeWarning("Halted")
        running_step = self.steps[self.current_step]
        step_quantum_state = self.compose_quantum_state(running_step)
        if not self.current_quantum_state:
            self.current_quantum_state = step_quantum_state
        else:
            self.current_quantum_state = step_quantum_state | self.current_quantum_state
        self.current_step += 1

    def run(self):
        for _ in self.steps:
            self.step()
        self.current_state = self.HALT_STATE

    def reset(self):
        self.current_step = 0
        self.current_quantum_state = None
        self.current_state = self.HALT_STATE

    def print(self):
        print("=" * 3 * len(self.steps))
        for line_no in range(self.n_qubits):
            line = ''
            for step in self.steps:
                state_repr = repr(step[line_no])
                line += state_repr
            print(line)
        print("=" * 3 * len(self.steps))


class QuantumProcessor(object):
    HALT_STATE = "HALT"
    RUNNING_STATE = "RUNNING"

    def __init__(self, circuit: QuantumCircuit):
        if circuit.rows != circuit.n_qubits:
            raise Exception("Declared circuit with n_qubits: {} but called add_row: {}".format(
                circuit.n_qubits, circuit.rows
            ))
        self.circuit = circuit
        self.c_step = 0
        self.c_q_state = None
        self.c_state = self.HALT_STATE
        self.reset()

    def compose_quantum_state(self, step):
        if any([type(s) is TwoQubitPartial for s in step]):
            two_qubit_gates = filter(lambda s: type(s) is TwoQubitPartial, step)
            state = []
            for two_qubit_gate in two_qubit_gates:
                two_qubit_gate.operator.verify_step(step)
                state = two_qubit_gate.operator.compose(step, state)
            return state
        return reduce((lambda x, y: x * y), step)

    def step(self):
        if self.c_step == 0 and self.c_state == self.HALT_STATE:
            self.c_state = self.RUNNING_STATE
        if self.c_step >= len(self.circuit.steps):
            self.c_state = self.HALT_STATE
            raise RuntimeWarning("Halted")
        step_quantum_state = self.get_next_step()
        if not self.c_q_state:
            self.c_q_state = step_quantum_state
        else:
            self.c_q_state = State((step_quantum_state | self.c_q_state).m)
        self.c_step += 1

    def get_next_step(self):
        running_step = self.circuit.steps[self.c_step]
        step_quantum_state = self.compose_quantum_state(running_step)
        return step_quantum_state

    def run(self):
        for _ in self.circuit.steps:
            self.step()
        self.c_state = self.HALT_STATE

    def reset(self):
        self.c_step = 0
        self.c_q_state = None
        self.c_state = self.HALT_STATE

    def measure(self):
        if self.c_state != self.HALT_STATE:
            raise RuntimeError("Processor is still running")
        return self.c_q_state.measure()

    def run_n(self, n: int):
        for i in range(n):
            self.run()
            result = self.measure()
            print("Run {}: {}".format(i, result))
            self.reset()

    def get_sample(self, n: int):
        rv = defaultdict(int)
        for i in range(n):
            self.run()
            result = self.measure()
            rv[result] += 1
            self.reset()
        return rv

    @staticmethod
    def print_sample(rv):
        for k, v in sorted(rv.items(), key=lambda x: x[0]):
            print("{}: {}".format(k, v))


def test_quantum_processor():
    # Produce Bell state between 0 and 2 qubit
    qc = QuantumCircuit(3)
    qc.add_row([H, C_partial])
    qc.add_row([_, _])
    qc.add_row([_, x_partial])
    qc.print()
    qp = QuantumProcessor(qc)
    qp.get_sample(100)


if __name__ == "__main__":
    test()
    # test_quantum_processor()