import random
from collections import defaultdict
from functools import reduce
from numbers import Complex
from pprint import pprint
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):
        """
        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_prob = 1.0 - sum(weights)
        format_str = self.get_fmt_of_element()
        choices = [REPR_EMPTY_SET] + [format_str.format(i) for i in range(len(weights))]
        weights = [empty_prob] + 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()] += 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 test_krisi_measurement_2():
    from qiskit import (
        QuantumCircuit,
        execute,
        Aer)
    simulator = Aer.get_backend('qasm_simulator')

    CASE_IDENTICAL = "identical"
    CASE_ORTHOGONAL = "orthogonal"

    def perform_exp(case):
        # produce angles with uniform distribution on the sphere
        t = round(np.random.uniform(0, 1), 10)
        theta0 = np.arccos(1 - 2 * t)
        phi0 = round(np.random.uniform(0, 2 * np.pi), 10)

        # rotate the 0th qubit in (theta, phi)
        q0 = State.from_bloch_angles(theta0, phi0)

        # rotate the 1st qubit depending on the case we are exploring...
        if case == CASE_IDENTICAL:
            # ... for identical we are rotating the 1st qubit the same angles as
            # the 0th
            theta1, phi1 = theta0, phi0
        else:
            # for orthogonal we rotate the 1st qubit in 90 degrees
            # orthogonal to the first
            theta1, phi1 = theta0 + np.pi, phi0

        q1 = State.from_bloch_angles(theta1, phi1)

        # Measure in the Bell's basis
        st = State(q0 * q1)
        meas = st.measure(basis=[b_phi_p, b_phi_m, b_psi_p, b_psi_m])
        return meas

    correct = 0
    cases = defaultdict(int)
    results = defaultdict(int)
    for i in range(1000):
        case = random.choice([CASE_IDENTICAL, CASE_ORTHOGONAL])
        cases[case] += 1
        result = perform_exp(case=case)
        results[result] += 1
        if result == '11':
            guess = CASE_ORTHOGONAL
            assert guess == case
        else:
            guess = CASE_IDENTICAL
        if guess == case:
            correct += 1
    print("Correct: {}".format(correct))
    # print("Results: {}".format(results))
    # print("Cases: {}".format(cases))


def test_krisi_measurement_3():
    beta = 2.6
    b_0 = State(1 / np.sqrt(2) * (s("|100>") - s("|010>")))
    b_1 = State(1 / np.sqrt(2) * (s("|011>") - s("|101>")))
    meas = s("|000>").measure(basis=[
        s("|000>"),
        State(np.cos(beta) * b_0 + np.sin(beta) * s("|001>")),
        State(-np.sin(beta) * b_0 + np.cos(beta) * s("|001>")),
        b_0,
        b_1,
        State(np.cos(beta) * b_1 + np.sin(beta) * s("|001>")),
        State(-np.sin(beta) * b_1 + np.cos(beta) * s("|001>")),
        s("|111>"),
    ])
    print(meas)


if __name__ == "__main__":
    test()
    test_quantum_processor()
    test_light()
    test_krisi_measurement_2()
    # TODO:
    test_krisi_measurement_3()