"""
Block vector and matrices.

Intended to be used with theano symbolic tensors


Also includes the BT "pseude namespace". It wraps the theano.tensor and theano.tensor.nnet
namespace in a way that the (elementwise) functions can be applied to block structures.
"""
import scipy as sp
import theano.tensor as T
from numbers import Number


class EmptyTensor:
    singleton = None

    def __new__(cls):
        if cls.singleton is None:
            cls.singleton = super().__new__(cls)
        return cls.singleton


class BlockApplication:
    """
    Encapsulate a block operation.

    This could have been a closure, but is
    implemented as class for its ``__repr__`` method.
    """
    def __init__(self, operation, name):
        self.operation = operation
        self.name = name

    def __repr__(self):
        return "<BlockOperation {}>".format(self.name)

    def __call__(self, block_item):
        return block_item.apply_op(self.operation)


class _BlockTensorOperation:
    """
    Helper to for support of theano style function application
    """
    def __init__(self, namespace):
        self.namespace = namespace

    def __getattr__(self, item):
        operation = getattr(self.namespace, item)
        apply = BlockApplication(operation, item)
        return apply


class _BlockTensorModuleList:
    _main = _BlockTensorOperation(T)
    nnet = _BlockTensorOperation(T.nnet)

    def __getattr__(self, item):
        return getattr(self._main, item)


# this is a hack to user theano functions on block vectors
BT = _BlockTensorModuleList()


class DimensionMismatch(Exception):
    pass


class BlockBase:
    def __init__(self, components):
        self._components = list(components)

    def __iter__(self):
        return iter(self._components)

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

    def eval(self):
        return type(self)(entry.eval() if hasattr(entry, "eval") else entry for entry in self)

    def apply_op(self, operation):
        return type(self)(operation(component) for component in self)

    @property
    def dtype(self):
        return [entry.dtype for entry in self]

    def __mul__(self, other):
        if isinstance(other, Number):
            return type(self)(component * other for component in self)
        else:
            return NotImplemented

    def __rmul__(self, other):
        if isinstance(other, Number):
            return self * other

    def __add__(self, other):
        if type(other) != type(self):
            return NotImplemented
        if len(self) != len(other):
            raise DimensionMismatch("{} != {}".format(len(self), len(other)))
        return type(self)(s + o for s, o in zip(self, other))

    def __eq__(self, other):
        try:
            if len(self) != len(other):
                return False
            for s, o in zip(self, other):
                res = s == o
                try:
                    res = res.all()
                except AttributeError:
                    pass
                if not res:
                    return False
            return True
        except TypeError:
            return


def inner_product(a, b):
    if isinstance(a, Number) and a == 0:
        return 0
    if isinstance(b, Number) and b == 0:
        return 0
    if isinstance(a, Number) and isinstance(b, Number):
        return a * b
    return a.dot(b)


class BlockVector(BlockBase):
    def __repr__(self):
        return "{}({})".format(self.__class__.__name__, self._components)
    
    def __format__(self, format_spec):
        if format_spec == "short":
            return str(self._components)
        else:
            return str(self)
    
    def __getitem__(self, item):
        if isinstance(item, slice):
            return type(self)(self._components[item])
        else:
            return self._components[item]

    @property
    def shape(self):
        return len(self._components),

    @property
    def inner_shapes(self):
        return type(self)(c.shape if hasattr(c, "shape") else "shape({})".format(c) for c in self._components)

    def __setitem__(self, key, value):
        self._components[key] = value
    
    def dot(self, other):
        return sum(inner_product(s, o) for s, o in zip(self, other))

    def to_dense_array(self):
        return sp.concatenate(self)

    @property
    def T(self):
        return type(self)(c.T for c in self)


class BlockMatrix(BlockBase):
    def __init__(self, components):
        super().__init__(BlockVector(component) for component in components)

    def __repr__(self):
        clsname = self.__class__.__name__
        components = "[" + format(self._components[0], "short")
        for c in self._components[1:]:
            components += ",\n" + " "*13 +format(c, "short")
        components += "]"
        return clsname + "(" + components + ")"

    def apply_op(self, operation):
        return type(self)(operation(c) for component in self for c in component)

    @property
    def T(self):
        block_transpose = zip(*self)
        transpoe = [[col.T if not isinstance(col, Number) else col for col in row] for row in block_transpose]
        return type(self)(transpoe)

    @property
    def shape(self):
        return len(self), len(self._components[0])

    @property
    def inner_shapes(self):
        return type(self)(c.inner_shapes for c in self._components)

    def __getitem__(self, item):
        if isinstance(item, tuple) and len(item) == 2:
            rows, cols = item
            if isinstance(rows, int) and isinstance(cols, int):
                return self._components[rows][cols]
            if isinstance(rows, int) and isinstance(cols, slice):
                r = self._components[rows]
                return type(r)(r[cols])
            if isinstance(rows, slice) and isinstance(cols, int):
                cls = type(self._components[0])
                return cls(c[cols] for c in self)
            if isinstance(rows, slice) and isinstance(cols, slice):
                return type(self)(type(row)(row[cols]) for row in self._components[rows])
            else:
                raise NotImplementedError()
        else:
            return self._components[item]

    def __setitem__(self, key, value):
        if isinstance(key, tuple):
            row, col = key
            self._components[row][col] = value
        else:
            raise NotImplementedError("Item assignment only implemented for integer indices")

    def dot(self, other):
        if isinstance(other, BlockVector):
            return type(other)(inner_product(s, other) for s in self)
        else:
            raise NotImplementedError()

    def flatten(self):
        cls = type(self._components[0])
        return cls(entry for row in self for entry in row)

    def __matmul__(self, other):
        if isinstance(other, BlockVector):
            return self.dot(other)
        else:
            return NotImplemented

    def __rmatmul__(self, other):
        if not isinstance(other, BlockVector):
            return NotImplemented
        else:
            return self.T @ other

    def to_dense_array(self):
        return sp.concatenate([component.to_dense_array() for component in self], axis=1)

    def get_diag(self, n=0):
        min_shape = min(self.shape) - n
        cls = type(self[0])
        return cls(self._components[k][k+n] for k in range(min_shape))

    def set_diag(self, values, n=0):
        min_shape = min(self.shape) - n
        for k, value in zip(range(min_shape), values):
            self._components[k][k+n] = value

    @classmethod
    def from_list(cls, list, nr_rows):
        nr_cols = len(list) //nr_rows
        return cls(list[start*nr_cols:start*nr_cols+nr_cols] for start in range(len(list)//nr_cols))
