from __future__ import annotations

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

import numpy as np
from numbers import Complex

from load_test import sizeof_fmt

ListOrNdarray = Union[list, np.ndarray]


class Matrix(object):
    """Represents a quantum state as a vector in Hilbert space"""

    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 other.__class__(self.m + other.m)

    def __eq__(self, other):
        if isinstance(other, Complex):
            return bool(np.allclose(self.m, other))
        elif isinstance(other, TwoQubitGate):
            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 self.__class__(self.m * other)
        elif isinstance(other, Matrix):
            return other.__class__(np.kron(self.m, other.m))
        raise NotImplementedError

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

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

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

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

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

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

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

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

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

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


class State(Matrix):
    def __init__(self, m: ListOrNdarray = None, name: str = '', *args, **kwargs):
        """An Operator turns one function into another"""
        super().__init__(m)
        self.name = name

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

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

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

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

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

    def measure(self):
        weights = [self.get_prob(j) for j in range(len(self))]
        format_str = "{:0" + str(int(np.ceil(np.log2(len(weights))))) + "b}"
        choices = [format_str.format(i) for i in range(len(weights))]
        return random.choices(choices, weights)[0]


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

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

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


class UnitaryMatrix(Matrix):
    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 is
          1. square and
          2. the product of itself with conjugate transpose is Identity UU+ = I"""
        is_square = self.m.shape[0] == self.m.shape[1]
        UU_ = np.dot(self._conjugate_transpose(), self.m)
        I = np.eye(self.m.shape[0])
        return is_square and np.isclose(UU_, I).all()


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

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


"""
Define States and Operators 
"""

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

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

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

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

CNOT = UnitaryOperator([[1, 0, 0, 0],
                        [0, 1, 0, 0],
                        [0, 0, 0, 1],
                        [0, 0, 1, 0], ])


# TODO - 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 TwoQubitGate(object):
    def __init__(self, rpr):
        self.rpr = rpr

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


C = TwoQubitGate("C")
x = TwoQubitGate("x")

# TODO: End Hacky way to define 2-qbit gate
###########################################################


def test():
    _p = State([[1 / np.sqrt(2)], [1 / np.sqrt(2)]], name='+')

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

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

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

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

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

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

    # III: Measurement | A quantum measurement is described by an orthonormal basis |e_j>
    #                    for state space. If the initial state of the system is |psi>
    #                    then we get outcome j with probability pr(j) = |<e_j|psi>|^2

    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

    # IV: Compositing | tensor/kronecker product when composing
    assert _0 * _0 == State([[1], [0], [0], [0]])
    assert _0 * _1 == State([[0], [1], [0], [0]])
    assert _1 * _0 == State([[0], [0], [1], [0]])
    assert _1 * _1 == State([[0], [0], [0], [1]])

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

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

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

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

    ################################
    # TODO: Don't know where outer product fits...
    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 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 QuantumProcessor(object):
    HALT_STATE = "HALT"
    RUNNING_STATE = "RUNNING"

    def __init__(self, n_qubits: int):
        self.n_qubits = n_qubits
        self.steps = [[_0 for _ in range(n_qubits)], ]
        self._called_add_row = 0
        self.current_step = 0
        self.current_quantum_state = None
        self.current_state = self.HALT_STATE
        self.reset()

    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):
        if C in step or x in step:
            if not (step.count(C) == 1 and step.count(x) == 1):
                raise RuntimeError("Both CONTROL and TARGET need to be defined in the same step exactly once")
            # TODO: Hacky way to do CNOT
            #       Should generalize for a 2-Qubit gate
            # _0.x(_0) * Matrix(I.m) + _1.x(_1) * Matrix(X.m)
            outer_0, outer_1 = [], []
            for s in step:
                if s == C:
                    outer_0.append(_0.x(_0))
                    outer_1.append(_1.x(_1))
                elif s == x:
                    outer_0.append(Matrix(I.m))
                    outer_1.append(Matrix(X.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)
        return reduce((lambda x, y: x * y), step)

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

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

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

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

    def measure(self):
        if self.current_state != self.HALT_STATE:
            raise RuntimeError("Processor is still running")
        return self.current_quantum_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()
        for k, v in sorted(rv.items(), key=lambda x: x[0]):
            print("{}: {}".format(k, v))
        return rv


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


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