from .function import *
from .helper import *
from tortto import np

"""
contain functions that can't be generated by grad_fcn_generator 
"""


class Split(Function):  # keep input _version: True
    @staticmethod
    def forward(ctx, *inputs, **params):
        xt0, = inputs
        x0 = xt0.data
        dim = params['dim']
        split_size_or_sections = params['split_size_or_sections']
        dim_size = x0.shape[dim]
        if dim < 0:
            dim += x0.ndim
        if split_size_or_sections.__class__ is int:
            split_size = split_size_or_sections
            ytn = tuple(
                build_links(
                    x0[
                        tuple(
                            slice(None) if i != dim else slice(j, j + split_size) for i in range(x0.ndim)
                        )
                    ],
                    grad_fn=ctx,
                    _output_idx=j // split_size
                )
                for j in range(0, dim_size, split_size)
            )
        else:
            sections = split_size_or_sections
            if sum(sections) != dim_size:
                raise RuntimeError(f"split_with_sizes expects split_sizes to sum exactly to {dim_size} "
                                   f"(input tensor's size at dimension {dim}), but got split_sizes={sections}")
            sum_sections = np.cumsum(split_size_or_sections)
            ytn = tuple(
                build_links(
                    x0[
                        tuple(
                            slice(None) if i != dim else slice(sum_sections[j] - sec, sum_sections[j]) for i in
                            range(x0.ndim)
                        )
                    ],
                    grad_fn=ctx,
                    _output_idx=j
                )
                for j, sec in enumerate(sections)
            )
        ctx.save_for_backward(xt0)
        ctx.params['output_shapes'] = tuple(yt.shape for yt in ytn)
        return ytn

    @staticmethod
    def backward(ctx, *grad_outputs):
        x0, = ctx.saved_tensors
        dim = ctx.params['dim']
        output_shapes = ctx.params['output_shapes']
        xp = ctx.xp
        grad0 = xp.concatenate(
            [
                xp.zeros(output_shapes[i], dtype=x0.dtype) if gn is None else gn
                for i, gn in enumerate(grad_outputs)
            ],
            axis=dim
        )
        return grad0


class Expand(Function):  # keep input _version: True
    @staticmethod
    def forward(ctx, *inputs, **params):
        xt0, = inputs
        x0 = xt0.data
        sizes = params['sizes']
        xp = ctx.xp
        leading_dims = len(sizes) - len(x0.shape)
        strides = [0] * leading_dims + list(x0.strides)
        x0_singleton_dims = []  # singleton axes to be summed during backward
        for i in range(len(sizes)):
            if i < leading_dims:  # leading dimensions
                if sizes[i] <= 0:
                    raise RuntimeError(f"The expanded size of the tensor ({sizes[i]}) isn't allowed in a leading, "
                                       f"non-existing dimension {i}")
            else:
                i -= len(sizes)  # for non-leading dimensions, count backward
                if x0.shape[i] == 1:
                    if sizes[i] > 1:
                        x0_singleton_dims.append(i)
                        strides[i] = 0
                else:
                    if sizes[i] != -1 and x0.shape[i] != sizes[i]:
                        raise RuntimeError(f"The expanded size of the tensor ({sizes[i]}) must match the existing size "
                                           f"({x0.shape[i]}) at non-singleton dimension {i + len(sizes)}.  "
                                           f"Target sizes: {sizes}.  Tensor sizes: {x0.shape}")
        y0 = xp.lib.stride_tricks.as_strided(x0, shape=sizes, strides=strides)  # a numpy/cupy array
        yt0 = build_links(y0, grad_fn=ctx)  # convert to a new nparray/cparray, version is 0
        yt0.data._version = x0._version  # keep version
        ctx.params = {'x0_singleton_dims': x0_singleton_dims, 'leading_dims': leading_dims}
        return yt0

    @staticmethod
    def backward(ctx, *grad_outputs):
        g0, = grad_outputs
        x0_singleton_dims = tuple(ctx.params['x0_singleton_dims'])
        leading_dims = tuple(range(ctx.params['leading_dims']))
        grad0 = g0.sum(x0_singleton_dims + leading_dims, keepdims=True).squeeze(leading_dims)
        return grad0
