from numbers import Integral

import numpy as np
from numpy.lib.stride_tricks import as_strided


def sliding_window_view(arr, window_shape, step, dilation=None):
    """Create a sliding window view over the trailing dimensions of an array.
    No copy is made unless the input array is not contiguous in memory.

    The window is applied only to valid regions of ``arr``, but is applied greedily.

    See Notes section for details.

    Parameters
    ----------
    arr : numpy.ndarray, shape=(..., [x, (...), z])
        C-contiguous array over which sliding view-window is applied along the trailing
        dimensions ``[x, ..., z]``, as determined by the length of ``window_shape``.

        If ``arr`` is not C-contiguous, it will be replaced by ``numpy.ascontiguousarray(arr)``

    window_shape : Sequence[int]
        Specifies the shape of the view-window: ``[Wx, (...), Wz]``.
        The length of `window_shape` determines the length of ``[x, (...) , z]``.

    step : Union[int, Sequence[int]]
        The step sized used along the ``[x, (...), z]`` dimensions: ``[Sx, (...), Sz]``.
        If a single integer is specified, a uniform step size is used.

    dilation : Optional[Union[int, Sequence[int]]]
        The dilation factor used along the ``[x, (...), z]`` directions: ``[Dx, (...), Dz]``.
        If no value is specified, a dilation factor of 1 is used along each direction.
        Dilation specifies the step size used when filling the window's elements

    Returns
    -------
    numpy.ndarray
        A contiguous view of ``arr``, of shape ``([X, (...), Z], ..., [Wx, (...), Wz])``, where
        ``[X, ..., Z]`` is the shape of the grid on which the window was applied. See Notes
        sections for more details.

    Raises
    ------
    ValueError, TypeError
        Invalid step-size, window shape, or dilation

    Notes
    -----
    Window placement:
        Given a dimension of size x, with a window of size W along this dimension, applied
        with stride S and dilation D, the window will be applied::
                                  X = (x - (W - 1) * D + 1) // S + 1
        number of times along that dimension.

    Interpreting output:
        In general, given an array ``arr`` of shape (..., x, (...), z), and::

            out = sliding_window_view(arr, window_shape=[Wx, (...), Wz], step=[Sx, (...), Sz])

        then indexing ``out`` with ``[ix, (...), iz]`` produces the following view of ``x``::

            out[ix, (...), iz] ==
                x[..., ix*Sx:(ix*Sx + Wx*Dx):Dx, (...), iz*Sz:(iz*Sz + Wz*Dz):Dz]

        For example, suppose ``arr`` is an array of shape-(10, 12, 6). Specifying sliding
        window of shape ``(3, 3)`` with step size ``(2, 2)``, dilation ``(2, 1)`` will create the view::

                        [[arr[:,  0:6:2, 0:3], arr[:,   0:6:3, 3:6]]
                         [arr[:, 6:12:2, 0:3], arr[:, 6:12:12, 3:6]]]

        producing a view of shape ``(2, 2, 10, 3, 3)`` in total.

    Examples
    --------
    >>> import numpy as np
    >>> x = np.arange(36).reshape(6, 6)
    >>> x
    array([[ 0,  1,  2,  3,  4,  5],
           [ 6,  7,  8,  9, 10, 11],
           [12, 13, 14, 15, 16, 17],
           [18, 19, 20, 21, 22, 23],
           [24, 25, 26, 27, 28, 29],
           [30, 31, 32, 33, 34, 35]])

    Apply an 3x2 window with step-sizes of (2, 2). This results in
    the window being placed twice along axis-0 and three times along axis-1.

    >>> y = sliding_window_view(x, step=(2, 2), window_shape=(3, 2))
    >>> y.shape
    (2, 3, 3, 2)

    window applied at (0, 0)

    >>> y[0, 0]
    array([[ 0,  1],
           [ 6,  7],
           [12, 13]])

    window applied at (2, 0)

    >>> y[1, 0]
    array([[12, 13],
           [18, 19],
           [24, 25]])

    window applied at (0, 2)

    >>> y[0, 1]
    array([[ 2,  3],
           [ 8,  9],
           [14, 15]])

    verify that an element in this window-view is correct

    >>> i, j = np.random.randint(0, 2, size=2)
    >>> wx, wy = (2, 2)
    >>> sx, sy = (2, 2)
    >>> np.all(y[i, j] == x[..., i*sx:(i*sx + wx), j*sy:(j*sy + wy)])
    True
    """

    if not hasattr(window_shape, "__iter__"):
        raise TypeError(
            f"`window_shape` must be a sequence of positive integers, got: {window_shape}"
        )
    window_shape = tuple(window_shape)
    if not all(isinstance(i, Integral) and i > 0 for i in window_shape):
        raise TypeError(
            f"`window_shape` must be a sequence of positive integers, "
            f"got: {window_shape}"
        )

    if len(window_shape) > arr.ndim:
        raise ValueError(
            f"`window_shape` ({window_shape}) cannot specify more values than "
            f"`arr.ndim` ({arr.ndim})."
        )

    if not isinstance(step, Integral) and not hasattr(step, "__iter__"):
        raise TypeError(
            f"`step` must be a positive integer or a sequence of positive "
            f"integers, got: {step}"
        )

    step = (
        (int(step),) * len(window_shape) if isinstance(step, Integral) else tuple(step)
    )

    if not all(isinstance(i, Integral) and i > 0 for i in step):
        raise ValueError(
            f"`step` must be a positive integer or a sequence of positive "
            f"integers, got: {step}"
        )

    if any(i > j for i, j in zip(window_shape[::-1], arr.shape[::-1])):
        raise ValueError(
            f"Each size of the window-shape must fit within the trailing "
            f"dimensions of `arr`."
            f"{window_shape} does not fit in {arr.shape[-len(window_shape) :]}"
        )

    if (
        dilation is not None
        and not isinstance(dilation, Integral)
        and not hasattr(dilation, "__iter__")
    ):
        raise TypeError(
            f"`dilation` must be None, a positive integer, or a sequence of "
            f"positive integers, got: {dilation}"
        )
    if dilation is None:
        dilation = np.ones((len(window_shape),), dtype=int)
    else:
        if isinstance(dilation, Integral):
            dilation = np.full((len(window_shape),), fill_value=dilation, dtype=int)
        else:
            np.asarray(dilation)

        if not all(isinstance(i, Integral) and i > 0 for i in dilation) or len(
            dilation
        ) != len(window_shape):
            raise ValueError(
                f"`dilation` must be None, a positive integer, or a sequence of "
                f"positive integers with the same length as `window_shape` "
                f"({window_shape}), got: {dilation}"
            )
        if any(
            w * d > s
            for w, d, s in zip(window_shape[::-1], dilation[::-1], arr.shape[::-1])
        ):
            raise ValueError(
                f"The dilated window ({tuple(w * d for w, d in zip(window_shape, dilation))}) "
                f"must fit within the trailing "
                f"dimensions of `arr` ({arr.shape[-len(window_shape) :]})"
            )

    if not arr.flags["C_CONTIGUOUS"]:
        arr = np.ascontiguousarray(arr)

    step = np.array(step)  # (Sx, ..., Sz)
    window_shape = np.array(window_shape)  # (Wx, ..., Wz)
    in_shape = np.array(arr.shape[-len(step) :])  # (x, ... , z)
    nbyte = arr.strides[-1]  # size, in bytes, of element in `arr`

    # per-byte strides required to fill a window
    win_stride = tuple(np.cumprod(arr.shape[:0:-1])[::-1]) + (1,)

    # per-byte strides required to advance the window
    step_stride = tuple(win_stride[-len(step) :] * step)

    # update win_stride to accommodate dilation
    win_stride = np.array(win_stride)
    win_stride[-len(step) :] *= dilation
    win_stride = tuple(win_stride)

    # tuple of bytes to step to traverse corresponding dimensions of view
    # see: 'internal memory layout of an ndarray'
    stride = tuple(int(nbyte * i) for i in step_stride + win_stride)

    # number of window placements along x-dim: X = (x - (Wx - 1)*Dx + 1) // Sx + 1
    out_shape = tuple((in_shape - ((window_shape - 1) * dilation + 1)) // step + 1)

    # ([X, (...), Z], ..., [Wx, (...), Wz])
    out_shape = out_shape + arr.shape[: -len(step)] + tuple(window_shape)
    out_shape = tuple(int(i) for i in out_shape)

    return as_strided(arr, shape=out_shape, strides=stride, writeable=False)
