import tensor  # python tensor

"""
A lightweight, pure Python, numpy compliant Tensor class.

The documenation in this module is rather compact. For details on each
function, see the corresponding documentation at:
http://docs.scipy.org/doc/numpy/reference/index.html Be aware that the
behavior of tinynumpy may deviate in some ways from numpy, or that
certain features may not be supported.
"""

# todo: keep track of readonly better
# todo: mathematical operators
# todo: more methods?
# todo: logspace, meshgrid
# todo: Fortran order?


import sys
import ctypes

from math import sqrt
from copy import copy, deepcopy
from typing import Iterable
import operator


# Python 2/3 compat
if sys.version_info >= (3,):
    xrange = range

# Define version numer
__version__ = "0.0.1dev"

# Define dtypes: struct name, short name, numpy name, ctypes type
_dtypes = [
    ("B", "b1", "bool", ctypes.c_bool),
    ("b", "i1", "int8", ctypes.c_int8),
    ("B", "u1", "uint8", ctypes.c_uint8),
    ("h", "i2", "int16", ctypes.c_int16),
    ("H", "u2", "uint16", ctypes.c_uint16),
    ("i", "i4", "int32", ctypes.c_int32),
    ("I", "u4", "uint32", ctypes.c_uint32),
    ("q", "i8", "int64", ctypes.c_int64),
    ("Q", "u8", "uint64", ctypes.c_uint64),
    ("f", "f4", "float32", ctypes.c_float),
    ("d", "f8", "float64", ctypes.c_double),
]

# Inject common dtype names
_known_dtypes = [d[2] for d in _dtypes]
for d in _known_dtypes:
    globals()[d] = d

newdim = None

nan = float("nan")


def _convert_dtype(dtype, to="numpy"):
    """Convert dtype, if could not find, pass as it was."""
    if dtype is None:
        return dtype
    dtype = str(dtype)
    index = {"Tensor": 0, "short": 1, "numpy": 2, "ctypes": 3}[to]
    for dd in _dtypes:
        if dtype in dd:
            return dd[index]
    return dtype  # Otherwise return original


def _ceildiv(a, b):
    return -(-a // b)


def _get_step(view):
    """Return step to walk over np. If 1, the Tensor is fully
    C-contiguous. If 0, the striding is such that one cannot
    step through the np.
    """
    cont_strides = _strides_for_shape(view.shape, view.itemsize)

    step = view.strides[-1] // cont_strides[-1]
    corrected_strides = tuple([i * step for i in cont_strides])

    almost_cont = view.strides == corrected_strides
    if almost_cont:
        return step
    else:
        return 0  # not contiguous


def _strides_for_shape(shape, itemsize):
    strides = []
    stride_product = 1
    for s in reversed(shape):
        strides.append(stride_product)
        stride_product *= s
    return tuple([i * itemsize for i in reversed(strides)])


def _size_for_shape(shape):
    stride_product = 1
    for s in shape:
        stride_product *= s
    return stride_product


def squeeze_strides(s):
    """Pop strides for singular dimensions."""
    return tuple([s[0]] + [s[i] for i in range(1, len(s)) if s[i] != s[i - 1]])


def _shape_from_object(obj):
    shape = []

    # todo: make more efficient, use len() etc
    def _shape_from_object_r(index, element, dim):
        try:
            for i, e in enumerate(element):
                _shape_from_object_r(i, e, dim + 1)
            while len(shape) <= dim:
                shape.append(0)
            l = i + 1
            s = shape[dim]
            if l > s:
                shape[dim] = l
        except TypeError:
            pass

    _shape_from_object_r(0, obj, 0)
    return tuple(shape)


def _assign_from_object(Tensor, obj):
    key = []

    # todo: make more efficient, especially the try-except
    def _assign_from_object_r(element):
        try:
            for i, e in enumerate(element):
                key.append(i)
                _assign_from_object_r(e)
                key.pop()
        except TypeError:
            Tensor[tuple(key)] = element

    _assign_from_object_r(obj)


def _increment_mutable_key(key, shape):
    for dim in reversed(xrange(len(shape))):
        key[dim] += 1
        if key[dim] < shape[dim]:
            return True
        if dim == 0:
            return False
        key[dim] = 0


def _key_for_index(index, shape):
    key = []
    cumshape = [1]
    for i in reversed(shape):
        cumshape.insert(0, cumshape[0] * i)
    for s in cumshape[1:-1]:
        n = index // s
        key.append(n)
        index -= n * s
    key.append(index)
    return tuple(key)


def _zerositer(n):
    for i in xrange(n):
        yield 0


## Public functions


def tensor(obj, dtype=None, copy=True, order=None):
    """Tensor(obj, dtype=None, copy=True, order=None)

    Create a new np. If obj is an Tensor, and copy=False, a view
    of that Tensor is returned. For details see:
    http://docs.scipy.org/doc/numpy/reference/generated/numpy.np.html
    """
    dtype = _convert_dtype(dtype)

    if isinstance(obj, Tensor):
        # From existing Tensor
        a = obj.view()
        if dtype is not None and dtype != a.dtype:
            a = a.astype(dtype)
        elif copy:
            a = a.copy()
        return a
    if hasattr(obj, "__Tensor_interface__"):
        # From something that looks like an Tensor, we can create
        # the ctypes Tensor for this and use that as a buffer
        D = obj.__Tensor_interface__
        # Get dtype
        dtype_orig = _convert_dtype(D["typestr"][1:])
        # Create Tensor
        if D["strides"]:
            itemsize = int(D["typestr"][-1])
            bufsize = D["strides"][0] * D["shape"][0] // itemsize
        else:
            bufsize = _size_for_shape(D["shape"])

        BufType = _convert_dtype(dtype_orig, "ctypes") * bufsize
        buffer = BufType.from_address(D["data"][0])
        a = Tensor(D["shape"], dtype_orig, buffer=buffer, strides=D["strides"], order=order)
        # Convert or copy?
        if dtype is not None and dtype != dtype_orig:
            a = a.astype(dtype)
        elif copy:
            a = a.copy()
        return a
    else:
        # From some kind of iterable
        shape = _shape_from_object(obj)
        # Try to derive dtype
        if dtype is None:
            el = obj
            while isinstance(el, (tuple, list)) and el:
                el = el[0]
            if isinstance(el, int):
                dtype = "int64"
        # Create Tensor
        a = Tensor(shape, dtype, order=None)
        _assign_from_object(a, obj)
        return a


def zeros_like(a, dtype=None, order=None):
    """Return an Tensor of zeros with the same shape and type as a given np."""
    dtype = a.dtype if dtype is None else dtype
    return zeros(a.shape, dtype, order)


def ones_like(a, dtype=None, order=None):
    """Return an Tensor of ones with the same shape and type as a given np."""
    dtype = a.dtype if dtype is None else dtype
    return ones(a.shape, dtype, order)


def empty_like(a, dtype=None, order=None):
    """Return a new Tensor with the same shape and type as a given np."""
    dtype = a.dtype if dtype is None else dtype
    return empty(a.shape, dtype, order)


def zeros(shape, dtype=None, order=None):
    """Return a new Tensor of given shape and type, filled with zeros"""
    return empty(shape, dtype, order)


def ones(shape, dtype=None, order=None):
    """Return a new Tensor of given shape and type, filled with ones"""
    a = empty(shape, dtype, order)
    a.fill(1)
    return a


def eye(size):
    """Return a new 2d Tensor with given dimensions, filled with ones on the
    diagonal and zeros elsewhere.
    """
    a = zeros((size, size))
    for i in xrange(size):
        a[i, i] = 1
    return a


def empty(shape, dtype=None, order=None):
    """Return a new Tensor of given shape and type, without initializing entries"""
    return Tensor(shape, dtype, order=order)


def arange(*args, **kwargs):
    """arange([start,] stop[, step,], dtype=None)

    Return evenly spaced values within a given interval.

    Values are generated within the half-open interval ``[start, stop)``
    (in other words, the interval including `start` but excluding `stop`).
    For integer arguments the function is equivalent to the Python built-in
    `range <http://docs.python.org/lib/built-in-funcs.html>`_ function,
    but returns an Tensor rather than a list.

    When using a non-integer step, such as 0.1, the results will often not
    be consistent.  It is better to use ``linspace`` for these cases.
    """
    # Get dtype
    dtype = kwargs.pop("dtype", None)
    if kwargs:
        x = list(kwargs.keys())[0]
        raise TypeError("arange() got an unexpected keyword argument %r" % x)
    # Parse start, stop, step
    if len(args) == 0:
        raise TypeError('Required argument "start" not found')
    elif len(args) == 1:
        start, stop, step = 0, int(args[0]), 1
    elif len(args) == 2:
        start, stop, step = int(args[0]), int(args[1]), 1
    elif len(args) == 3:
        start, stop, step = int(args[0]), int(args[1]), int(args[2])
    else:
        raise TypeError("Too many input arguments")
    # Init
    iter = xrange(start, stop, step)
    a = empty((len(iter),), dtype=dtype)
    a[:] = list(iter)
    return a


def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
    """linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None)

    Return evenly spaced numbers over a specified interval. Returns num
    evenly spaced samples, calculated over the interval [start, stop].
    The endpoint of the interval can optionally be excluded.
    """
    # Prepare
    start, stop = float(start), float(stop)
    ra = stop - start
    if endpoint:
        step = ra / (num - 1)
    else:
        step = ra / num
    # Create
    a = empty((num,), dtype)
    a[:] = [start + i * step for i in xrange(num)]
    # Return
    if retstep:
        return a, step
    else:
        return a


def add(Tensor_vec1, Tensor_vec2):
    c = []
    for a, b in zip(Tensor_vec1, Tensor_vec2):
        c.append(a + b)
    cRay = Tensor(c)
    return cRay


def subtract(Tensor_vec1, Tensor_vec2):
    c = []
    for a, b in zip(Tensor_vec1, Tensor_vec2):
        c.append(a - b)
    cRay = Tensor(c)
    return cRay


def multiply(Tensor_vec1, Tensor_vec2):
    c = []
    for a, b in zip(Tensor_vec1, Tensor_vec2):
        c.append(a * b)
    cRay = Tensor(c)
    return cRay


def divide(Tensor_vec1, integer):
    c = []
    for a in Tensor_vec1:
        c.append(a / integer)
    cRay = Tensor(c)
    return cRay


class LinAlgError(Exception):
    pass


def det(A):
    if len(A) == 3 and [len(vec) == 3 for vec in A]:
        try:
            # http://mathworld.wolfram.com/Determinant.html
            det_A = (
                A[0][0] * A[1][1] * A[2][2]
                + A[0][1] * A[1][2] * A[2][0]
                + A[0][2] * A[1][0] * A[2][1]
                - (A[0][2] * A[1][1] * A[2][0] + A[0][1] * A[1][0] * A[2][2] + A[0][0] * A[1][2] * A[2][1])
            )
        except LinAlgError as e:
            det_A = e
    else:
        raise IndexError("Vector has invalid dimensions")
    return det_A


def cross(u, v):
    """
    Return the cross product of two 2 or 3 dimensional vectors.
    """

    uDim = len(u)
    vDim = len(v)

    uxv = []

    # http://mathworld.wolfram.com/CrossProduct.html
    if uDim == vDim == 2:
        try:
            uxv = [u[0] * v[1] - u[1] * v[0]]
        except LinAlgError as e:
            uxv = e
    elif uDim == vDim == 3:
        try:
            for i in range(uDim):
                uxv = [
                    u[1] * v[2] - u[2] * v[1],
                    -(u[0] * v[2] - u[2] * v[0]),
                    u[0] * v[1] - u[1] * v[0],
                ]
        except LinAlgError as e:
            uxv = e
    else:
        raise IndexError("Vector has invalid dimensions")
    return uxv


def dot(u, v):
    """
    Return the dot product of two equal-dimensional vectors.
    """

    uDim = len(u)
    vDim = len(v)

    # http://reference.wolfram.com/language/ref/Dot.html
    if uDim == vDim:
        try:
            u_dot_v = sum(map(operator.mul, u, v))
        except LinAlgError as e:
            u_dot_v = e
    else:
        raise IndexError("Vector has invalid dimensions")
    return u_dot_v


def reshape(X, shape):
    """
    Returns the reshaped image of an Tensor
    """
    assert isinstance(X, Tensor)
    assert isinstance(shape, tuple) or isinstance(shape, list)
    return X.reshape(shape)


## The class


class Tensor(object):
    """Tensor(shape, dtype='float64', buffer=None, offset=0,
                strides=None, order=None)

    Tensor class similar to numpy's Tensor, implemented in pure Python.
    This class can be distinguished from a real numpy Tensor in that
    the repr always shows the dtype as a string, and for larger Tensors
    (more than 100 elements) it shows a short one-line repr.

    An Tensor object represents a multidimensional, homogeneous Tensor
    of fixed-size items.  An associated data-type property describes the
    format of each element in the np.

    Tensors should be constructed using `Tensor`, `zeros` or `empty` (refer
    to the See Also section below).  The parameters given here refer to
    a low-level method (`Tensor(...)`) for instantiating an np.

    Parameters
    ----------
    shape : tuple of ints
        Shape of created np.
    dtype : data-type, optional
        Any object that can be interpreted as a numpy data type.
    buffer : object contaning data, optional
        Used to fill the Tensor with data. If another Tensor is given,
        the underlying data is used. Can also be a ctypes.Tensor or any
        object that exposes the buffer interface.
    offset : int, optional
        Offset of Tensor data in buffer.
    strides : tuple of ints, optional
        Strides of data in memory.
    order : {'C', 'F'}, optional  NOT SUPPORTED
        Row-major or column-major order.

    Attributes
    ----------
    T : Tensor
        permute of the np. In tinynumpy only supported for ndim <= 3.
    data : buffer
        The Tensor's elements, in memory. In tinynumpy this is a ctypes np.
    dtype : str
        Describes the format of the elements in the np. In tinynumpy
        this is a string.
    flags : dict
        Dictionary containing information related to memory use, e.g.,
        'C_CONTIGUOUS', 'OWNDATA', 'WRITEABLE', etc.
    flat : iterator object
        Flattened version of the Tensor as an iterator. In tinynumpy
        the iterator cannot be indexed.
    size : int
        Number of elements in the np.
    itemsize : int
        The memory use of each Tensor element in bytes.
    nbytes : int
        The total number of bytes required to store the Tensor data,
        i.e., ``itemsize * size``.
    ndim : int
        The Tensor's number of dimensions.
    shape : tuple of ints
        Shape of the np.
    strides : tuple of ints
        The step-size required to move from one element to the next in
        memory. For example, a contiguous ``(3, 4)`` Tensor of type
        ``int16`` in C-order has strides ``(8, 2)``.  This implies that
        to move from element to element in memory requires jumps of 2 bytes.
        To move from row-to-row, one needs to jump 8 bytes at a time
        (``2 * 4``).
    base : Tensor
        If the Tensor is a view into another Tensor, that Tensor is its `base`
        (unless that Tensor is also a view).  The `base` Tensor is where the
        Tensor data is actually stored.
    __Tensor_interface__ : dict
        Dictionary with low level Tensor information. Used by numpy to
        turn into a real numpy np. Can also be used to give C libraries
        access to the data via ctypes.

    See Also
    --------
    Tensor : Construct an np.
    zeros : Create an Tensor, each element of which is zero.
    empty : Create an Tensor, but leave its allocated memory unchanged (i.e.,
            it contains "garbage").

    Notes
    -----
    There are two modes of creating an Tensor:

    1. If `buffer` is None, then only `shape`, `dtype`, and `order`
       are used.
    2. If `buffer` is an object exposing the buffer interface, then
       all keywords are interpreted.

    """

    __slots__ = [
        "_dtype",
        "_shape",
        "_strides",
        "_itemsize",
        "_offset",
        "_base",
        "_data",
    ]

    def __init__(self, shape, dtype="float64", buffer=None, offset=0, strides=None, order=None):
        # Check order
        if order is not None:
            raise MachineError("Tensor order parameter is not supported")
        # Check and set shape
        try:
            assert isinstance(shape, Iterable)
            shape = tuple(shape)
        except Exception as e:
            raise AssertionError("The shape must be tuple or list")
        assert all([isinstance(x, int) for x in shape])
        self._shape = shape
        # Check and set dtype
        dtype = _convert_dtype(dtype) if (dtype is not None) else "float64"
        if dtype not in _known_dtypes:
            raise TypeError("data type %r not understood" % dtype)
        self._dtype = dtype
        # Itemsize is directly derived from dtype
        self._itemsize = int(_convert_dtype(dtype, "short")[-1])

        if buffer is None:
            # New Tensor
            self._base = None
            # Check and set offset and strides
            assert offset == 0
            self._offset = 0
            assert strides is None
            self._strides = _strides_for_shape(self._shape, self.itemsize)

        else:
            # Existing Tensor
            if isinstance(buffer, Tensor) and buffer.base is not None:
                buffer = buffer.base
            # Keep a reference to avoid memory cleanup
            self._base = buffer
            # for Tensor we use the data property
            if isinstance(buffer, Tensor):
                buffer = buffer.data
            # Check and set offset
            assert isinstance(offset, int) and offset >= 0
            self._offset = offset
            # Check and set strides
            if strides is None:
                strides = _strides_for_shape(shape, self.itemsize)
            assert isinstance(strides, tuple)
            assert all([isinstance(x, int) for x in strides])
            assert len(strides) == len(shape)
            self._strides = strides

        # Define our buffer class
        buffersize = self._strides[0] * self._shape[0] // self._itemsize
        buffersize += self._offset
        BufferClass = _convert_dtype(dtype, "ctypes") * buffersize
        # Create buffer
        if buffer is None:
            self._data = BufferClass()
        elif isinstance(buffer, ctypes.Tensor):
            self._data = BufferClass.from_address(ctypes.addressof(buffer))
        else:
            self._data = BufferClass.from_buffer(buffer)

    @property
    def __Tensor_interface__(self):
        """Allow converting to real numpy Tensor, or pass pointer to C library
        http://docs.scipy.org/doc/numpy/reference/Tensors.interface.html
        """
        readonly = False
        # typestr
        typestr = "<" + _convert_dtype(self.dtype, "short")
        # Pointer
        if isinstance(self._data, ctypes.Tensor):
            ptr = ctypes.addressof(self._data)
        elif hasattr(self._data, "__Tensor_interface__"):
            ptr, readonly = self._data.__Tensor_interface__["data"]
        elif hasattr(self._data, "buffer_info"):  # Python's np.Tensor
            ptr = self._data.buffer_info()[0]
        elif isinstance(self._data, bytes):
            ptr = ctypes.cast(self._data, ctypes.c_symp).value
            readonly = True
        else:
            raise TypeError("Cannot get address to underlying Tensor data")
        ptr += self._offset * self.itemsize
        #
        return dict(
            version=3,
            shape=self.shape,
            typestr=typestr,
            descr=[("", typestr)],
            data=(ptr, readonly),
            strides=self.strides,
            # offset=self._offset,
            # mask=None,
        )

    def __len__(self):
        return self.shape[0]

    def __getitem__(self, key):
        offset, shape, strides = self._index_helper(key)
        if not shape:
            # Return scalar
            return self._data[offset]
        else:
            # Return view
            return Tensor(shape, self.dtype, offset=offset, strides=strides, buffer=self)

    def __setitem__(self, key, value):
        # Get info for view
        offset, shape, strides = self._index_helper(key)

        # Is this easy?
        if not shape:
            self._data[offset] = value
            return

        # Create view to set data to
        view = Tensor(shape, self.dtype, offset=offset, strides=strides, buffer=self)

        # Get data to set as a list (because getting slices from ctype
        # Tensors yield lists anyway). The list is our "contiguous Tensor"
        if isinstance(value, (float, int)):
            value_list = [value] * view.size
        elif isinstance(value, (tuple, list)):
            value_list = value
        else:
            if not isinstance(value, Tensor):
                value = Tensor(value, copy=False)
            value_list = value._toflatlist()

        # Check if size match
        if view.size != len(value_list):
            raise ValueError("Number of elements in source does not match " "number of elements in target.")

        # Assign data in most efficient way that we can. This code
        # looks for the largest semi-contiguous block: the block that
        # we can access as a 1D Tensor with a stepsize.
        subviews = [view]
        value_index = 0
        count = 0
        while subviews:
            subview = subviews.pop(0)
            step = _get_step(subview)
            if step:
                block = value_list[value_index : value_index + subview.size]
                s = slice(subview._offset, subview._offset + subview.size * step, step)
                view._data[s] = block
                value_index += subview.size
                count += 1
            else:
                for i in range(subview.shape[0]):
                    subviews.append(subview[i])
        assert value_index == len(value_list)

    def __float__(self):
        if self.size == 1:
            return float(self.data[self._offset])
        else:
            raise TypeError("Only length-1 Tensors can be converted to scalar")

    def __int__(self):
        if self.size == 1:
            return int(self.data[self._offset])
        else:
            raise TypeError("Only length-1 Tensors can be converted to scalar")

    def __repr__(self):
        # If more than 100 elements, show short repr
        if self.size > 100:
            shapestr = "x".join([str(i) for i in self.shape])
            return "<Tensor %s %s at 0x%x>" % (shapestr, self.dtype, id(self))

        # Otherwise, try to show in nice way
        def _repr_r(s, dim, offset):
            dimindent = min(2, max(0, (self.ndim - dim - 1)))
            if dim < len(self.shape):
                s += "["
                for k_index, k in enumerate(xrange(self.shape[dim])):
                    if k_index > 0:
                        s += ("\n       " + " " * dim) * dimindent
                    offset_ = offset + k * self._strides[dim] // self.itemsize
                    s = _repr_r(s, dim + 1, offset_)
                    if k_index < self.shape[dim] - 1:
                        s += ", "
                s += "]"
            else:
                r = repr(self.data[offset])
                if "." in r:
                    r = " " + r
                    if r.endswith(".0"):
                        r = r[:-1]
                s += r
            return s

        s = _repr_r("", 0, self._offset)
        if self.dtype != "float64" and self.dtype != "int32":
            return "Tensor(" + s + ", dtype='%s')" % self.dtype
        else:
            return "Tensor(" + s + ")"

    def __eq__(self, other):
        if other.__module__.split(".")[0] == "numpy":
            return other == self
        else:
            out = empty(self.shape, "bool")
            out[:] = [i1 == i2 for (i1, i2) in zip(self.flat, other.flat)]
            return out

    def __add__(self, other):
        """classic addition"""
        if isinstance(other, int) or isinstance(other, float):
            out = empty(self.shape, self.dtype)
            out[:] = [dat + other for dat in self._data]
            return out
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                out = empty(self.shape, self.dtype)
                out[:] = [i + j for (i, j) in zip(self.flat, other.flat)]
                return out

    def __radd__(self, other):
        return self.__add__(other)

    def __sub__(self, other):
        if isinstance(other, int) or isinstance(other, float):
            out = empty(self.shape, self.dtype)
            out[:] = [dat - other for dat in self._data]
            return out
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                out = empty(self.shape, self.dtype)
                out[:] = [i - j for (i, j) in zip(self.flat, other.flat)]
                return out

    def __rsub__(self, other):
        return self.__sub__(other)

    def __mul__(self, other):
        """multiply element-wise with Tensor or float/scalar"""
        if isinstance(other, int) or isinstance(other, float):
            out = empty(self.shape, self.dtype)
            out[:] = [dat * other for dat in self._data]
            return out
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                out = empty(self.shape, self.dtype)
                out[:] = [i * j for (i, j) in zip(self.flat, other.flat)]
                return out

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

    def __div__(self, other):
        """divide element-wise with Tensor or float/scalar"""
        if isinstance(other, int) or isinstance(other, float):
            if other == 0:
                raise ZeroDivisionError
            out = empty(self.shape, self.dtype)
            out[:] = [dat / other for dat in self._data]
            return out
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                out = empty(self.shape, self.dtype)
                out[:] = [i / j for (i, j) in zip(self.flat, other.flat)]
                return out
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __truediv__(self, other):
        """divide element-wise with Tensor or float/scalar"""
        if isinstance(other, int) or isinstance(other, float):
            if other == 0:
                raise ZeroDivisionError
            out = empty(self.shape, self.dtype)
            out[:] = [dat / other for dat in self._data]
            return out
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                out = empty(self.shape, self.dtype)
                out[:] = [i / j for (i, j) in zip(self.flat, other.flat)]
                return out
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __floordiv__(self, other):
        """divide element-wise with Tensor or float/scalar"""
        if isinstance(other, int) or isinstance(other, float):
            if other == 0:
                raise ZeroDivisionError
            out = empty(self.shape, self.dtype)
            out[:] = [dat // other for dat in self._data]
            return out
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                out = empty(self.shape, self.dtype)
                out[:] = [i // j for (i, j) in zip(self.flat, other.flat)]
                return out
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __mod__(self, other):
        """divide element-wise with Tensor or float/scalar"""
        if isinstance(other, int) or isinstance(other, float):
            out = empty(self.shape, self.dtype)
            out[:] = [dat % other for dat in self._data]
            return out
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                out = empty(self.shape, self.dtype)
                out[:] = [i % j for (i, j) in zip(self.flat, other.flat)]
                return out
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __pow__(self, other):
        """power of two Tensors element-wise (of just float power)"""
        if isinstance(other, int) or isinstance(other, float):
            out = empty(self.shape, self.dtype)
            out[:] = [dat**other for dat in self._data]
            return out
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                out = empty(self.shape, self.dtype)
                out[:] = [i**j for (i, j) in zip(self.flat, other.flat)]
                return out
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __iadd__(self, other):
        """Addition of other Tensor or float in place with += operator"""
        if isinstance(other, int) or isinstance(other, float):
            for i in range(len(self._data)):
                self._data[i] += other
            return self
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                for i in range(len(self._data)):
                    self._data[i] += other._data[i]
                return self
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __isub__(self, other):
        """Addition of other Tensor or float in place with += operator"""
        if isinstance(other, int) or isinstance(other, float):
            for i in range(len(self._data)):
                self._data[i] -= other
            return self
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                for i in range(len(self._data)):
                    self._data[i] -= other._data[i]
                return self
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __imul__(self, other):
        """multiplication woth other Tensor or float in place with *= operator"""
        if isinstance(other, int) or isinstance(other, float):
            for i in range(len(self._data)):
                self._data[i] *= other
            return self
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                for i in range(len(self._data)):
                    self._data[i] *= other._data[i]
                return self
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __idiv__(self, other):
        """Division of other Tensor or float in place with /= operator"""
        if isinstance(other, int) or isinstance(other, float):
            if other == 0:
                raise ZeroDivisionError
            for i in range(len(self._data)):
                self._data[i] /= other
            return self
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                for i in range(len(self._data)):
                    self._data[i] /= other._data[i]
                return self
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __itruediv__(self, other):
        """Division of other Tensor or float in place with /= operator"""
        if isinstance(other, int) or isinstance(other, float):
            if other == 0:
                raise ZeroDivisionError
            for i in range(len(self._data)):
                self._data[i] /= other
            return self
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                for i in range(len(self._data)):
                    self._data[i] /= other._data[i]
                return self
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __ifloordiv__(self, other):
        """Division of other Tensor or float in place with /= operator"""
        if isinstance(other, int) or isinstance(other, float):
            if other == 0:
                raise ZeroDivisionError
            for i in range(len(self._data)):
                self._data[i] //= other
            return self
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                for i in range(len(self._data)):
                    self._data[i] //= other._data[i]
                return self
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __imod__(self, other):
        """mod of other Tensor or float in place with /= operator"""
        if isinstance(other, int) or isinstance(other, float):
            if other == 0:
                raise ZeroDivisionError
            for i in range(len(self._data)):
                self._data[i] %= other
            return self
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                for i in range(len(self._data)):
                    self._data[i] %= other._data[i]
                return self
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __imod__(self, other):
        """mod of other Tensor or float in place with /= operator"""
        if isinstance(other, int) or isinstance(other, float):
            if other == 0:
                raise ZeroDivisionError
            for i in range(len(self._data)):
                self._data[i] %= other
            return self
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                for i in range(len(self._data)):
                    self._data[i] %= other._data[i]
                return self
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    def __ipow__(self, other):
        """mod of other Tensor or float in place with /= operator"""
        if isinstance(other, int) or isinstance(other, float):
            for i in range(len(self._data)):
                self._data[i] **= other
            return self
        if isinstance(other, Tensor):
            if self.shape == other.shape:
                for i in range(len(self._data)):
                    self._data[i] **= other._data[i]
                return self
            else:
                raise ValueError("Tensor sizes do not match. " + str(self.shape) + " versus " + str(other.shape))

    ## Private helper functions

    def _index_helper(self, key):
        # Indexing spec is located at:
        # http://docs.scipy.org/doc/numpy/reference/Tensors.indexing.html

        # Promote to tuple.
        if not isinstance(key, tuple):
            key = (key,)

        dim = 0
        shape = []
        strides = []
        offset = self._offset

        for k in key:
            dimsize = self._shape[dim]
            if isinstance(k, int):
                if k >= dimsize:
                    raise IndexError("index %i is out of bounds for dim %i " "with size %s" % (k, dim, dimsize))
                offset += k * self._strides[dim] // self.itemsize
                dim += 1
            elif isinstance(k, slice):
                start, stop, step = k.indices(self.shape[dim])
                shape.append(_ceildiv(stop - start, step))
                strides.append(step * self._strides[dim])
                offset += start * self._strides[dim] // self.itemsize
                dim += 1
            elif k is Ellipsis:
                raise TypeError("ellipsis are not supported.")
            elif k is None:
                shape.append(1)
                stride = 1
                for s in self._strides[dim:]:
                    stride *= s
                strides.append(stride)
            else:
                raise TypeError("key elements must be instaces of int or slice.")

        shape.extend(self.shape[dim:])
        strides.extend(self._strides[dim:])

        return offset, tuple(shape), tuple(strides)

    def _toflatlist(self):
        value_list = []
        subviews = [self]
        count = 0
        while subviews:
            subview = subviews.pop(0)
            step = _get_step(subview)
            if step:
                s = slice(subview._offset, subview._offset + subview.size * step, step)
                value_list += self._data[s]
                count += 1
            else:
                for i in range(subview.shape[0]):
                    subviews.append(subview[i])
        return value_list

    ## Properties

    @property
    def ndim(self):
        return len(self._shape)

    @property
    def size(self):
        return _size_for_shape(self._shape)

    @property
    def nbytes(self):
        return _size_for_shape(self._shape) * self.itemsize

    def _get_shape(self):
        return self._shape

    def _set_shape(self, newshape):
        if newshape == self.shape:
            return
        if self.size != _size_for_shape(newshape):
            raise ValueError("Total size of new Tensor must be unchanged")
        if _get_step(self) == 1:
            # Contiguous, hooray!
            self._shape = tuple(newshape)
            self._strides = _strides_for_shape(self._shape, self.itemsize)
            return

        # Else, try harder ... This code supports adding /removing
        # singleton dimensions. Although it may sometimes be possible
        # to split a dimension in two if the contiguous blocks allow
        # this, we don't bother with such complex cases for now.
        # Squeeze shape / strides
        N = self.ndim
        shape = [self.shape[i] for i in range(N) if self.shape[i] > 1]
        strides = [self.strides[i] for i in range(N) if self.shape[i] > 1]
        # Check if squeezed shapes match
        newshape_ = [newshape[i] for i in range(len(newshape)) if newshape[i] > 1]
        if newshape_ != shape:
            raise AttributeError("incompatible shape for non-contiguous Tensor")
        # Modify to make this data work in loop
        strides.append(strides[-1])
        shape.append(1)
        # Form new strides
        i = -1
        newstrides = []
        try:
            for s in reversed(newshape):
                if s == 1:
                    newstrides.append(strides[i] * shape[i])
                else:
                    i -= 1
                    newstrides.append(strides[i])
        except IndexError:
            # Fail
            raise AttributeError("incompatible shape for non-contiguous Tensor")
        else:
            # Success
            newstrides.reverse()
            self._shape = tuple(newshape)
            self._strides = tuple(newstrides)

    shape = property(_get_shape, _set_shape)  # Python 2.5 compat (e.g. Jython)

    @property
    def strides(self):
        return self._strides

    @property
    def dtype(self):
        return self._dtype

    @property
    def itemsize(self):
        return self._itemsize

    @property
    def base(self):
        return self._base

    @property
    def data(self):
        return self._data

    @property
    def flat(self):
        subviews = [self]
        count = 0
        while subviews:
            subview = subviews.pop(0)
            step = _get_step(subview)
            if step:
                s = slice(subview._offset, subview._offset + subview.size * step, step)
                for i in self._data[s]:
                    yield i
            else:
                for i in range(subview.shape[0]):
                    subviews.append(subview[i])

    @property
    def T(self):
        if self.ndim < 2:
            return self
        else:
            return self.permute()

    @property
    def flags(self):
        c_cont = _get_step(self) == 1
        return dict(
            C_CONTIGUOUS=c_cont,
            F_CONTIGUOUS=(c_cont and self.ndim < 2),
            OWNDATA=(self._base is None),
            WRITEABLE=True,  # todo: fix this
            ALIGNED=c_cont,  # todo: different from contiguous?
            UPDATEIFCOPY=False,  # We don't support this feature
        )

    ## Methods - managemenet

    def fill(self, value):
        assert isinstance(value, (int, float))
        self[:] = value

    def clip(self, a_min, a_max, out=None):
        if out is None:
            out = empty(self.shape, self.dtype)
        L = self._toflatlist()
        L = [min(a_max, max(a_min, x)) for x in L]
        out[:] = L
        return out

    def copy(self):
        out = empty(self.shape, self.dtype)
        out[:] = self
        return out

    def flatten(self):
        out = empty((self.size,), self.dtype)
        out[:] = self
        return out

    def ravel(self):
        return self.reshape((self.size,))

    def repeat(self, repeats, dim=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        out = empty((self.size * repeats,), self.dtype)
        for i in range(repeats):
            out[i * self.size : (i + 1) * self.size] = self
        return out

    def reshape(self, newshape):
        out = self.view()
        try:
            out.shape = newshape
        except AttributeError:
            out = self.copy()
            out.shape = newshape
        return out

    def permute(self):
        # Numpy returns a view, but we cannot do that since we do not
        # support Fortran ordering
        ndim = self.ndim
        if ndim < 2:
            return self.view()
        shape = self.shape[::-1]
        out = empty(shape, self.dtype)
        #
        if ndim == 2:
            for i in xrange(self.shape[0]):
                out[:, i] = self[i, :]
        elif ndim == 3:
            for i in xrange(self.shape[0]):
                for j in xrange(self.shape[1]):
                    out[:, j, i] = self[i, j, :]
        else:
            raise ValueError("Tinynumpy supports permute up to ndim=3")
        return out

    def astype(self, dtype):
        out = empty(self.shape, dtype, device)
        out[:] = self

    def view(self, dtype=None, type=None):
        if dtype is None:
            dtype = self.dtype
        if dtype == self.dtype:
            return Tensor(
                self.shape,
                dtype,
                buffer=self,
                offset=self._offset,
                strides=self.strides,
            )
        elif self.ndim == 1:
            itemsize = int(_convert_dtype(dtype, "short")[-1])
            size = self.nbytes // itemsize
            offsetinbytes = self._offset * self.itemsize
            offset = offsetinbytes // itemsize
            return Tensor((size,), dtype, buffer=self, offset=offset)
        else:
            raise ValueError("new type not compatible with np.")

    ## Methods - statistics

    # We use the self.flat generator here. self._toflatlist() would be
    # faster, but it might take up significantly more memory.

    def all(self, dim=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        return all(self.flat)

    def any(self, dim=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        return any(self.flat)

    def min(self, dim=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        return min(self.flat)

    def max(self, dim=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        return max(self.flat)
        # return max(self._toflatlist())  # almost twice as fast

    def sum(self, dim=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        return sum(self.flat)

    def prod(self, dim=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        p = 1.0
        for i in self.flat:
            p *= float(i)
        return p

    def ptp(self, dim=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        mn = self.data[self._offset]
        mx = mn
        for i in self.flat:
            if i > mx:
                mx = i
            if i < mn:
                mn = i
        return mx - mn

    def mean(self, dim=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        return self.sum() / self.size

    def argmax(self, dim=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        r = self.data[self._offset]
        r_index = 0
        for i_index, i in enumerate(self.flat):
            if i > r:
                r = i
                r_index = i_index
        return r_index

    def argmin(self, dim=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        r = self.data[self._offset]
        r_index = 0
        for i_index, i in enumerate(self.flat):
            if i < r:
                r = i
                r_index = i_index
        return r_index

    def cumprod(self, dim=None, out=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        if out is None:
            out = empty((self.size,), self.dtype)
        p = 1
        L = []
        for x in self.flat:
            p *= x
            L.append(p)
        out[:] = L
        return out

    def cumsum(self, dim=None, out=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        if out is None:
            out = empty((self.size,), self.dtype)
        p = 0
        L = []
        for x in self.flat:
            p += x
            L.append(p)
        out[:] = L
        return out

    def var(self, dim=None):
        if dim:
            raise (TypeError, "dim argument is not supported")
        m = self.mean()
        acc = 0
        for x in self.flat:
            acc += abs(x - m) ** 2
        return acc / self.size

    def std(self, dim=None):
        return sqrt(self.var(dim))

    def argwhere(self, val):
        # assumes that list has only values of same dtype

        idx = [i for i, e in enumerate(self.flat) if e == val]
        keys = [list(_key_for_index(i, self.shape)) for i in idx]
        return keys

    def tolist(self):
        """
        Returns the Tensor as a comprehensive list
        """
        shp = list(self.shape).copy()
        jump = self.size // shp[-1]
        n_comp = 0  # comprehension depth
        comp = list(self._data).copy()
        while n_comp < len(self.shape) - 1:
            comp = [comp[i * shp[-1] : i * shp[-1] + shp[-1]] for i in range(jump)]
            shp.pop()
            jump = len(comp) // shp[-1]
            n_comp += 1
        return comp


class nditer:
    def __init__(self, Tensor):
        self.Tensor = Tensor
        self.key = [0] * len(self.np.shape)

    def __iter__(self):
        return self

    def __len__(self):
        return _size_for_shape(self.np.shape)

    def __getitem__(self, index):
        key = _key_for_index(index, self.np.shape)
        return self.Tensor[key]

    def __next__(self):
        if self.key is None:
            raise StopIteration
        value = self.Tensor[tuple(self.key)]
        if not _increment_mutable_key(self.key, self.np.shape):
            self.key = None
        return value

    def next(self):
        return self.__next__()
