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_EMPTY_SET = "Ø"
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 __sub__(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, PartialQubit):
            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 Matrix(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 __len__(self):
        return len(self.m)

    def inner(self, other):
        """Define inner product - a.k.a dot product, scalar product - <0|0>"""
        return Matrix(np.dot(self._conjugate_transpose(), other.m))

    def dot(self, other):
        """Alias to inner"""
        return self.inner(other)

    def scalar(self, other):
        """Alias to inner"""
        return self.inner(other)

    def outer(self, other):
        """Define outer product |0><0|
        https://en.wikipedia.org/wiki/Outer_product

        Given two vectors u and v, their outer product u ⊗ v is defined
        as the m × n matrix A obtained by multiplying each element of u
        by each element of v

        The outer product u ⊗ v is equivalent to a matrix multiplication uvT,
        provided that u is represented as a m × 1 column vector and v as a
        n × 1 column vector (which makes vT a row vector).
        """
        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 kron(self, other):
        """Define kronicker product - |0> ⊗ |0> = |00>"""
        return Matrix(np.kron(self.m, other.m))

    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 HorizontalVector(Matrix):
    """Horizontal vector is basically a list"""

    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 len(self.m.shape) == 1


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


VectorOrList = Union[Vector, Matrix, ListOrNdarray]


class State(Vector):
    def __init__(self, m: VectorOrList = None, name: str = '', *args, **kwargs):
        """State vector representing quantum state"""
        if type(m) in [Vector, Matrix]:
            super().__init__(m.m, *args, **kwargs)
        else:
            super().__init__(m, *args, **kwargs)
        self.name = name
        self.measurement_result = None
        # TODO: SHOULD WE NORMALIZE?
        # 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)

    @classmethod
    def from_bloch_angles(cls, theta, phi):
        """Creates a state from angles in Bloch sphere"""
        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_bloch_angles(self):
        """Returns the angles of this state on the Bloch sphere"""
        if not self.m.shape == (2, 1):
            raise Exception("State needs to describe only 1 qubit on the bloch sphere (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 to_ket(self):
        return self.conjugate_transpose()

    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.inner(self)).m).item(0)

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

    def get_prob_from_measurement_op(self, m_op):
        assert type(m_op) is MeasurementOperator
        return m_op.get_prob(self)

    def get_computational_basis_vector(self, j):
        return Vector([[1] if i == int(j) else [0] for i in range(len(self.m))])

    def get_prob(self, j, basis=None):
        """pr(j) = |<e_j|self>|^2
        """
        if basis is None:
            basis = self.get_computational_basis_vector(j)
        m = MeasurementOperator.create_from_basis(basis)
        return m.get_prob(self)

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

    def measure(self, basis=None, allow_empty=False):
        """
        m_ops: a set of MeasurementOperators.
               e.g. for computational basis, the m_ops is [|0><0|, |1><1|]
        Measures in the computational basis.
        Irreversable operation. Measuring again will result in the same result
        TODO: Generalize the method so it takes a basis
        TODO: 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")
        """
        if self.measurement_result:
            return self.measurement_result
        weights, m_ops = [], []
        if not basis:
            for j in range(len(self)):
                # Default is computational basis
                e_j = self.get_computational_basis_vector(j)
                m = MeasurementOperator.create_from_basis(e_j)
                m_ops.append(m)
        else:
            assert len(basis) == len(self)
            for i, e_j in enumerate(basis):
                # All should be orthogonal
                if i != 0:
                    assert e_j.is_orthogonal(basis[0])
                m_ops.append(MeasurementOperator.create_from_basis(e_j))

        for m_op in m_ops:
            prob = self.get_prob_from_measurement_op(m_op)
            weights += [prob]

        empty_choices, empty_weights = [], []
        if allow_empty == True:
            empty_prob = 1.0 - sum(weights)
            empty_weights = [empty_prob]
            empty_choices = [REPR_EMPTY_SET]
        else:
            # normalize the weights
            weights = list(np.array(weights) / sum(weights))

        format_str = self.get_fmt_of_element()
        choices = empty_choices + [format_str.format(i) for i in range(len(weights))]
        weights = empty_weights + weights
        self.measurement_result = random.choices(choices, weights)[0]

        return self.measurement_result

    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).m / np.sqrt(mo.get_prob(self))
        return State(m)

    def measure_partial(self, qubit_n):
        """Partial measurement of state
        Measures the n-th qubit with probability sum(a_n),
        adjusting the probability of the rest of the state
        """
        max_qubits = int(np.log2(len(self)))
        if not (0 < qubit_n <= max_qubits):
            raise Exception("Partial measurement of qubit_n must be between 1 and {}".format(max_qubits))
        format_str = self.get_fmt_of_element()
        # e.g. for state |000>:
        # ['000', '001', '010', '011', '100', '101', '110', '111']
        bin_repr = [format_str.format(i) for i in range(len(self))]
        # e.g. for qbit_n = 2
        # [0, 0, 1, 1, 0, 0, 1, 1]
        partial_measurement_of_qbit = [int(b[qubit_n - 1]) for b in bin_repr]
        # [0, 1, 4, 5]
        weights, choices = defaultdict(list), defaultdict(list)
        for result in [1, 0]:
            indexes_for_p_0 = [i for i, index in enumerate(partial_measurement_of_qbit) if index == result]
            weights[result] = [self.get_prob(j) for j in indexes_for_p_0]
            choices[result] = [format_str.format(i) for i in indexes_for_p_0]
        weights_01 = [sum(weights[0]), sum(weights[1])]
        measurement_result = random.choices([0, 1], weights_01)[0]
        normalization_factor = np.sqrt(sum([np.abs(i) ** 2 for i in weights[measurement_result]]))
        new_m = weights[measurement_result] / normalization_factor
        return str(measurement_result), State(new_m.reshape((len(new_m), 1)))

    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_bloch_angles()
        x = np.sin(theta) * np.cos(phi)
        y = np.sin(theta) * np.sin(phi)
        z = np.cos(theta)
        return [x, y, z]


def test_measure_partial():
    state = b_phi_p
    state.measure_partial(1)


def normalize_state(vector: Vector):
    """Normalize a state by dividing by the square root of sum of the squares of states"""
    norm_coef = np.sqrt(np.sum(np.array(vector.m) ** 2))
    if norm_coef == 0:
        raise TypeError("zero-sum vector")
    return State(vector.m / norm_coef)


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


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: sp.Lambda, *args, **kwargs):
        """An Operator turns one function into another"""
        self.func = func

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

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


class LinearOperator(Operator):
    def __init__(self, func: sp.Lambda, *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):
        # A function is (jointly) linear in a given set of variables
        # if all second-order derivatives are identically zero (including mixed ones).
        # https://stackoverflow.com/questions/36283548/check-if-an-equation-is-linear-for-a-specific-set-of-variables
        expr, vars_ = self.func.expr, self.func.variables
        for x in vars_:
            for y in vars_:
                try:
                    if not sp.Eq(sp.diff(expr, x, y), 0):
                        return False
                except TypeError:
                    return False
        return True


def test_linear_operator():
    # Operators turn one vector into another
    # the times 2 operator should return the times two multiplication
    _x = sp.Symbol('x')
    _times_2 = LinearOperator(sp.Lambda(_x, 2 * _x))
    assert _times_2.on(5) == 10
    assert _times_2(5) == 10

    assert_raises(TypeError, "Not a linear operator", LinearOperator, sp.Lambda(_x, _x ** 2))


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 LinearTransformation(LinearOperator, Matrix):
    def __init__(self, m: ListOrNdarray, func=None, name: str = '', *args, **kwargs):
        """LinearTransformation (or linear map) inherits from both LinearOperator and a Matrix
        It is used to act on a State vector by defining the operator to be the dot product"""
        self.name = name
        Matrix.__init__(self, m=m, *args, **kwargs)
        self.func = func or self.operator_func
        LinearOperator.__init__(self, func=func, *args, **kwargs)

    def _is_linear(self):
        # Every matrix transformation is a linear transformation
        # https://www.mathbootcamps.com/proof-every-matrix-transformation-is-a-linear-transformation/
        return True

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

    def get_eigens(self):
        """ Returns (eigenvalue, eigenvector)
        M|v> = λ|v>  ->
          M   :is the operator (self)
          |v> :is eigenstate of M
          λ   :is the corresponding eigenvalue
        """
        eigenvalues, eigenvectors = np.linalg.eig(self.m)
        rv = []
        for i in range(0, len(eigenvectors)):
            eigenvalue = eigenvalues[i]
            eigenvector = HorizontalVector(eigenvectors[:, i])
            rv.append((eigenvalue, eigenvector))
        return rv


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

    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(LinearTransformation, HermitianMatrix):
    def __init__(self, m: ListOrNdarray, name: str = '', *args, **kwargs):
        """HermitianMatrix inherits from both LinearTransformation 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 Vector(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 PartialQubit(object):
    def __init__(self, rpr):
        self.rpr = rpr
        self.operator = None

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


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


class TwoQubitOperator(UnitaryOperator):
    def __init__(self, m: ListOrNdarray, A: PartialQubit, B: PartialQubit,
                 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 
"""

# EMPTY STATE
_E = State([[0],
            [0]],
           name=REPR_EMPTY_SET)

_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_phi_p = State([[1 / np.sqrt(2)],
                 [0],
                 [0],
                 [1 / np.sqrt(2)]],
                name="{}+".format(REPR_GREEK_PHI))

b_psi_p = State([[0],
                 [1 / np.sqrt(2)],
                 [1 / np.sqrt(2)],
                 [0]],
                name="{}+".format(REPR_GREEK_PSI))

b_phi_m = State([[1 / np.sqrt(2)],
                 [0],
                 [0],
                 [-1 / np.sqrt(2)]],
                name="{}-".format(REPR_GREEK_PHI))

b_psi_m = State([[0],
                 [1 / np.sqrt(2)],
                 [-1 / np.sqrt(2)],
                 [0]],
                name="{}-".format(REPR_GREEK_PSI))

well_known_states = [_p, _m, b_phi_p, b_psi_p, b_phi_m, b_psi_m]

_ = 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)


# TOFFOLLI_GATE = ThreeQubitOperator([
#     [1, 0, 0, 0, 0, 0, 0, 0],
#     [0, 1, 0, 0, 0, 0, 0, 0],
#     [0, 0, 1, 0, 0, 0, 0, 0],
#     [0, 0, 0, 1, 0, 0, 0, 0],
#     [0, 0, 0, 0, 1, 0, 0, 0],
#     [0, 0, 0, 0, 0, 1, 0, 0],
#     [0, 0, 0, 0, 0, 0, 0, 1],
#     [0, 0, 0, 0, 0, 0, 1, 0],
# ], C_partial, C_partial, x_partial, I, 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 MeasurementOperator(HermitianOperator):
    """Measurement operators are Hermitians: <ψ|M†_m M_m|ψ>"""

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

    def get_prob(self, state: State):
        """Returns result of <ψ|M†_m M_m|ψ>
        """
        # <ψ|
        state_ct = state.conjugate_transpose()
        #  This is: <ψ|   .   M†_m M_m . |ψ>
        return state_ct.m.dot(self.m.dot(state.m)).item().real


m0 = MeasurementOperator.create_from_basis(_0)
m1 = MeasurementOperator.create_from_basis(_1)


def test_measurement_ops():
    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_eigenstuff():
    assert LinearTransformation(m=[[1, 0], [0, 0]]).get_eigens() == \
           [(1.0, HorizontalVector([1., 0.])), (0., HorizontalVector([0., 1.]))]


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.inner(_0) == 1  # parallel have 1 product
    assert _0.inner(_1) == 0  # orthogonal have 0 product
    assert _0.is_orthogonal(_1)
    assert _1.inner(8 * _0) == 8 * _1.inner(_0)  # Inner product is linear multiplied by constants
    assert _p.inner(_1 + _0) == _p.inner(_1) + _p.inner(_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.inner(_1) == _1.inner(_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
    #
    test_linear_operator()

    test_eigenstuff()

    # 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.on(_1) == _0
    assert X.on(_0) == _1

    # Test the Y Pauli operator with complex number literal notation
    assert Y.on(_0) == State([[0],
                              [1j]])
    assert Y.on(_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

    assert _0.measure() == '0'
    assert _1.measure() == '1'

    assert s("|10>").measure() == '10'
    assert s("|10>").measure_partial(1) == ("1", _0)
    assert s("|10>").measure_partial(2) == ("0", _1)

    # measure in arbitrary basis
    _0.measure(basis=[_p, _m])
    s("|00>").measure(basis=[b_phi_p, b_phi_m, b_psi_m, b_psi_p])

    # Maximally entangled
    result, pms = b_phi_p.measure_partial(1)
    if result == "0":
        assert pms == _0
    elif result == "1":
        assert pms == _1

    # 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.on(_0)
    assert superpos == _p

    # Then CNOT the superposition with a |0> qubit
    bell = CNOT.on(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

    ################################
    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_bloch_angles()
        s = State.from_bloch_angles(*angles)
        assert q == s

    assert State.from_bloch_angles(0, 0) == _0
    assert State.from_bloch_angles(np.pi, 0) == _1
    assert State.from_bloch_angles(np.pi / 2, 0) == _p
    assert State.from_bloch_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_bloch_angles(theta, phi)
        assert s == qbit


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, PartialQubit)]
        # 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.inner(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 PartialQubit for s in step]):
            two_qubit_gates = filter(lambda s: type(s) is PartialQubit, 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.inner(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.print_sample(qp.get_sample(100))


def test_light():
    # http://alienryderflex.com/polarizer/
    hor_filter = MeasurementOperator.create_from_basis(Matrix(_0.m), name='h')
    diag_filter = MeasurementOperator.create_from_basis(Matrix(_p.m), name='d')
    ver_filter = MeasurementOperator.create_from_basis(Matrix(_1.m), name='v')

    def random_light():
        random_pol = Vector([[np.random.uniform(0, 1)], [np.random.uniform(0, 1)]])
        return normalize_state(random_pol)

    def experiment(id, random_ls, filters):
        total_light = defaultdict(int)
        human_state = "Light"
        for filter in filters:
            human_state += "->{}".format(filter.name)

        for r in random_ls:
            st = filters[0].on(r)
            for filter in filters:
                st = filter.on(st)
            st = State(st)
            total_light[st.measure(allow_empty=True)] += 1
        print("{}. {}:\n    {}".format(id, human_state, dict(total_light)))

    its = 100
    random_lights = [random_light() for _ in range(its)]

    # Just horizonal - should result in some light
    experiment(1, random_lights, [hor_filter, ])

    # Vertical after horizonal - should result in 0
    experiment(2, random_lights, [ver_filter, hor_filter])

    # TODO: Something is wrong here...
    # Vertical after diagonal after horizontal - should result in ~ 50% compared to only Horizontal
    experiment(3, random_lights, [ver_filter, diag_filter, hor_filter])


def generate_bins(count):
    format_str = "{:0" + str(int(np.ceil(np.log2(count)))) + "b}"
    return [format_str.format(i) for i in range(count)]


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