quantum/lib.py

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_TARGET = "⊕"
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))
        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)


MatrixOrList = Union[Matrix, ListOrNdarray]


class HorizontalVector(Matrix):
    """Horizontal vector is basically a list"""

    def __init__(self, m: MatrixOrList = None, *args, **kwargs):
        super().__init__(m, *args, **kwargs)
        if type(m) in [Matrix]:
            super().__init__(m.m, *args, **kwargs)
        else:
            super().__init__(m, *args, **kwargs)
        if not self._is_horizontal_vector():
            raise TypeError("Not a horizontal vector")

    def _is_horizontal_vector(self):
        return len(self.m.shape) == 1


class Vector(Matrix):
    def __init__(self, m: MatrixOrList = None, *args, **kwargs):
        super().__init__(m, *args, **kwargs)
        if type(m) in [Matrix]:
            super().__init__(m.m, *args, **kwargs)
        else:
            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, MatrixOrList]


class State(Vector):
    def __init__(self, m: VectorOrList = None, name: str = '',
                 allow_unnormalized=False, *args, **kwargs):
        """State vector representing quantum state"""
        if type(m) in [Vector, Matrix, State]:
            super().__init__(m.m, *args, **kwargs)
        else:
            super().__init__(m, *args, **kwargs)
        self.name = name
        self.measurement_result = None
        self.last_basis = None
        if not allow_unnormalized and not self._is_normalized():
            raise TypeError("Not a normalized state vector")

    def __hash__(self):
        return hash(str(self.m))

    def __le__(self, other):
        return repr(self) < repr(other)

    @staticmethod
    def normalize(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.abs(vector.m ** 2)))
        if norm_coef == 0:
            raise TypeError("zero-sum vector")
        new_m = vector.m / norm_coef
        return State(new_m)

    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)
        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=None, basis=None):
        """pr(j) = |<e_j|self>|^2
        """
        if basis is None:
            assert j is not 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: Should we memoize per basis?
              Currently it's memoized if the basis is the same
              If it's measured twice, should it return the same state?
              What about if measured twice but in different bases?
              E.g. measure1 -> computation -> A
                   measure2 -> +/- basis -> B
                   measure3 -> computation -> should it return A again or
                   random weighted?
                   measure4 -> +/- basis -> should it return B again or
                   random weighted?
        :return: binary representation of the measured qubit (e.g. "011")
        """
        weights, m_ops = [], []
        if not basis:
            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)
                basis.append(State(e_j))
        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))

        if self.measurement_result and self.last_basis == basis:
            return self.measurement_result

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

        empty_choices, empty_weights = [], []
        if allow_empty:
            empty_prob = 1.0 - sum(weights)
            empty_weights = [empty_prob]
            empty_choices = [REPR_EMPTY_SET]
        else:
            # TODO: This may be wrong
            # normalize the weights
            weights_sum = np.sum(weights)
            if not np.isclose(weights_sum, 1.0):
                weights = list(np.array(weights) / weights_sum)
            assert np.isclose(np.sum(weights), 1.0)

        weights = empty_weights + weights
        self.measurement_result = random.choices(empty_choices + basis, weights)[0]
        self.last_basis = basis
        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 State(s('|' + 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]


class Ket(State):
    def __init__(self, *args, **kwargs):
        """Alias: Ket is a state"""
        State.__init__(self, *args, **kwargs)


def s(q, name=None):
    """Helper method for creating state easily"""
    if type(q) == str:
        if q[0] == '|' and q[-1] == '>':
            # e.g. |000>
            state = State.from_string(q)
            state.name = name
            return state
        elif q[0] == '|' and q[-1] == '|' and "><" in q:
            # e.g. |0><0|
            f_qbit, s_qbit = q.split("<")[0], q.split(">")[1]
            if len(f_qbit) == len(s_qbit):
                s_qbit = '|' + s_qbit[1:-1] + '>'
                fs, ss = State.from_string(f_qbit), State.from_string(s_qbit)
                return DensityMatrix(fs.x(ss).m, name=name)
    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, **kwargs):
        return self(*args, **kwargs)

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


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 = '', partials=None,
                 decomposition=None, *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

        Partials - a list of partial operators that are used in a multi-qubit
        UnitaryOperator to define the operator on each qubit. Each element of
        the list is the Nth partial that is used, i.e. for the first - |0><0|,
        for the second - |1><1|
        """
        if np.shape(m) != (2, 2) and partials is None and decomposition is None:
            raise Exception("Please define partials or decomposition in "
                            "a non-single operator")
        UnitaryMatrix.__init__(self, m=m, *args, **kwargs)
        LinearTransformation.__init__(self, m=m, func=self.operator_func, *args,
                                      **kwargs)
        self.name = name
        self.partials = partials
        self.decomposition = decomposition

    def operator_func(self, other, which_qbit=None):
        this_cols, other_rows = np.shape(self.m)[1], np.shape(other.m)[0]
        if this_cols == other_rows and which_qbit is None:
            return State(np.dot(self.m, other.m))
        if which_qbit is None:
            raise Exception("Operating dim-{} operator on a dim-{} state. "
                            "Please specify which_qubit to operate on".format(
                this_cols, other_rows))
        total_qbits = int(np.log2(other_rows))
        if type(which_qbit) is list and len(which_qbit) == 1:
            which_qbit = which_qbit[0]
        if type(which_qbit) is int:
            # single qubit-gate
            assert this_cols == 2
            assert 0 <= which_qbit < total_qbits
            extended_m = [self if i == which_qbit else I for i in
                          range(total_qbits)]
            new_m = np.prod(extended_m).m
        elif type(which_qbit) is list:
            if self.decomposition:
                return self.decomposition(other, *which_qbit)
            else:
                # single or multiple qubit-gate
                assert 1 <= len(which_qbit) <= total_qbits
                assert len(which_qbit) == len(self.partials)
                assert all([q < total_qbits for q in which_qbit])
                this_gate_len = 2 ** (len(self.partials) - 1)
                bins = generate_bins(this_gate_len)
                extended_m, next_control = [[] for _ in bins], 0
                partial_mapping = dict(zip(which_qbit, self.partials))
                for qbit in range(total_qbits):
                    if qbit not in partial_mapping:
                        for i in range(this_gate_len):
                            extended_m[i].append(I)
                    else:
                        this_partial = partial_mapping[qbit]
                        # TODO: This works only with C_partial :(((
                        if this_partial == C_partial:
                            for i in range(this_gate_len):
                                bin_dig = bins[i][next_control]
                                extended_m[i].append(
                                    s("|" + bin_dig + "><" + bin_dig + "|"))
                            next_control += 1
                        else:
                            for i in range(this_gate_len - 1):
                                extended_m[i].append(I)
                            extended_m[-1].append(this_partial)
                new_m = sum([np.prod(e).m for e in extended_m])
        else:
            raise Exception(
                "which_qubit needs to be either an int of N-th qubit or list")

        return State(np.dot(new_m, other.m))

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


class Gate(UnitaryOperator):
    def __init__(self, *args, **kwargs):
        """Alias: Gates are Unitary Operators"""
        UnitaryOperator.__init__(self, *args, **kwargs)


class HermitianOperator(LinearTransformation, HermitianMatrix):
    def __init__(self, m: ListOrNdarray, name: str = '', *args, **kwargs):
        """HermitianOperator 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)


class DensityMatrix(HermitianOperator):
    def __init__(self, m: ListOrNdarray, name: str = '', *args, **kwargs):
        """The density matrix is a representation of a linear operator called
        the density operator. The density matrix is obtained from the density
        operator by choice of basis in the underlying space. In practice,
        the terms density matrix and density operator are often used
        interchangeably. Both matrix and operator are self-adjoint
        (or Hermitian), positive semi-definite, of trace one, and may
        be infinite-dimensional."""
        self.name = name
        HermitianOperator.__init__(self, m=m, *args, **kwargs)
        if not self._is_density():
            raise TypeError("Not a Density matrix")

    def _is_density(self):
        """Checks if it trace is 1"""
        return np.isclose(np.trace(self.m), 1.0)


"""
Define States and Operators 
"""

# EMPTY STATE
_E = Vector([[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))

# define common bases:
computational_basis = [_0, _1]
plus_minus_basis = [_p, _m]
bell_basis = [b_phi_p, b_psi_p, b_phi_m, b_psi_m]

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

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

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

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

Z = Gate([[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
#

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

# Rotation operators
# https://www.quantum-inspire.com/kbase/rx-gate/
Rx = lambda theta: Gate([[np.cos(theta / 2), -1j * np.sin(theta / 2)],
                         [-1j * np.sin(theta / 2), np.cos(theta / 2)]],
                        name="Rx")

Ry = lambda theta: Gate([[np.cos(theta / 2), -np.sin(theta / 2)],
                         [np.sin(theta / 2), np.cos(theta / 2)]],
                        name="Ry")

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

# See [T-Gate](https://www.quantum-inspire.com/kbase/t-gate/)
T = lambda phi: Gate([[1, 0],
                      [0, np.power(np.e, 1j * phi / 2)]],
                     name="T")

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

# Additionally for Toffoli: on a 4qubit system with C-CX (qbits 0,2,3)
# s("|0><0|") * I * s("|0><0|") * I \
# + s("|0><0|") * I * s("|1><1|") * I \
# + s("|1><1|") * I * s("|0><0|") * I \
# + s("|1><1|") * I * s("|1><1|") * X

C_partial = Gate(I.m, name="C")
x_partial = Gate(X.m, name=REPR_TARGET)
z_partial = Gate(Z.m, name=REPR_TARGET)

CNOT = Gate([
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 0, 1],
    [0, 0, 1, 0],
], name="CNOT", partials=[C_partial, x_partial])

CZ = Gate([
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 1, 0],
    [0, 0, 0, -1],
], name="CZ", partials=[C_partial, z_partial])

# Can't do partials for SWAP, but can be decomposed
# SWAP is decomposed into three CNOTs where the second one is flipped
SWAP = Gate([
    [1, 0, 0, 0],
    [0, 0, 1, 0],
    [0, 1, 0, 0],
    [0, 0, 0, 1],
], name="SWAP",
    decomposition=lambda state, x, y:
    CNOT.on(CNOT.on(CNOT.on(state, [x, y]), [y, x]), [x, y])
)

TOFF = Gate([
    [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],
], partials=[C_partial, C_partial, x_partial])


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_partials():
    # test single qbit state
    assert X.on(s("|000>"), which_qbit=1) == s("|010>")
    assert X.on(s("|000>"), which_qbit=[1]) == s("|010>")

    # normal 2 qbit state
    assert CNOT.on(s("|00>")) == s("|00>")
    assert CNOT.on(s("|10>")) == s("|11>")
    assert_raises(Exception,
                  "Operating dim-4 operator on a dim-8 state. Please specify "
                  "which_qubit to operate on",
                  CNOT.on, s("|100>"))

    # Test flipped CNOT
    assert CNOT.on(s("|01>"), which_qbit=[1, 0]) == s("|11>")
    assert CNOT.on(s("|001>"), which_qbit=[2, 0]) == s("|101>")

    # Test SWAP via 3 successive CNOTs, the second CNOT is flipped
    assert CNOT.on(CNOT.on(CNOT.on(s("|10>")), which_qbit=[1, 0])) == s("|01>")

    # Test SWAP via 3 successive CNOTs, the 0<->2 are swapped
    assert CNOT.on(CNOT.on(CNOT.on(s("|100>"), [0, 2]), [2, 0]), [0, 2]) == s(
        "|001>")

    # apply on 0, 1 of 3qbit state
    assert CNOT.on(s("|000>"), which_qbit=[0, 1]) == s("|000>")
    assert CNOT.on(s("|100>"), which_qbit=[0, 1]) == s("|110>")

    # apply on 1, 2 of 3qbit state
    assert CNOT.on(s("|000>"), which_qbit=[1, 2]) == s("|000>")
    assert CNOT.on(s("|010>"), which_qbit=[1, 2]) == s("|011>")

    # apply on 0, 2 of 3qbit state
    assert CNOT.on(s("|000>"), which_qbit=[0, 2]) == s("|000>")
    assert CNOT.on(s("|100>"), which_qbit=[0, 2]) == s("|101>")

    # apply on 0, 2 of 4qbit state
    assert CNOT.on(s("|1000>"), which_qbit=[0, 2]) == s("|1010>")

    # apply on 0, 3 of 4qbit state
    assert CNOT.on(s("|1000>"), which_qbit=[0, 3]) == s("|1001>")

    # test SWAP gate
    assert SWAP.on(s("|10>")) == s("|01>")
    assert SWAP.on(s("|01>")) == s("|10>")
    # Tests SWAP on far-away gates
    assert SWAP.on(s("|001>"), which_qbit=[0, 2]) == s("|100>")

    # test Toffoli gate
    assert TOFF.on(s("|000>")) == s("|000>")
    assert TOFF.on(s("|100>")) == s("|100>")
    assert TOFF.on(s("|110>")) == s("|111>")
    assert TOFF.on(s("|111>")) == s("|110>")

    # controls on 0, 1 target on 3 for 4
    assert TOFF.on(s("|1100>"), which_qbit=[0, 1, 3]) == s("|1101>")
    # controls on 0, 2 target on 3 for 4
    assert TOFF.on(s("|1010>"), which_qbit=[0, 2, 3]) == s("|1011>")
    assert TOFF.on(s("|1011>"), which_qbit=[0, 2, 3]) == s("|1010>")
    assert TOFF.on(s("|1111>"), which_qbit=[0, 2, 3]) == s("|1110>")
    # controls on 0, 2 target on 4 for 5
    assert TOFF.on(s("|10101>"), which_qbit=[0, 2, 4]) == s("|10100>")
    assert TOFF.on(s("|10111>"), which_qbit=[0, 2, 4]) == s("|10110>")
    assert TOFF.on(s("|11111>"), which_qbit=[0, 2, 4]) == s("|11110>")
    assert TOFF.on(s("|10100>"), which_qbit=[0, 2, 4]) == s("|10101>")


def test_sequential_measurements():
    st = H.on(s("|0>"))
    meas1 = st.measure()
    assert st.measurement_result is not None
    assert st.last_basis == computational_basis

    # Measuring more times should return the same result
    for i in range(10):
        meas2 = st.measure()
        assert meas2 == meas1

    # Measuring in different basis might not yield the same result
    meas3 = st.measure(basis=plus_minus_basis)
    assert st.measurement_result is not None
    assert st.last_basis == plus_minus_basis

    # But measuring more times should still return the same result
    for i in range(10):
        meas4 = st.measure(basis=plus_minus_basis)
        assert meas4 == meas3


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

    # Test 2+ qbit gates and partials
    test_partials()

    # 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() == s("|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()

    test_sequential_measurements()

    # 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
    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):
        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):
        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 rv.items():
            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 State.normalize(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, allow_unnormalized=True)
            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 get_bin_fmt(count):
    return "{:0" + str(int(np.ceil(np.log2(count)))) + "b}"


def generate_bin(i, count):
    format_str = get_bin_fmt(count)
    return format_str.format(i)


def generate_bins(count):
    format_str = get_bin_fmt(count)
    return [format_str.format(i) for i in range(count)]


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