import numpy as np

from .tensor import Function


# ************* unary ops *************

class ReLU(Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return np.maximum(input, 0)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        return grad_output * (input >= 0)


class Log(Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return np.log(input)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        return grad_output / input


class Exp(Function):
    @staticmethod
    def forward(ctx, input):
        ret = np.exp(input)
        ctx.save_for_backward(ret)
        return ret

    @staticmethod
    def backward(ctx, grad_output):
        ret, = ctx.saved_tensors
        return grad_output * ret


# ************* reduce ops *************

class Sum(Function):
    @staticmethod
    def forward(ctx, input, axis=None):
        ctx.save_for_backward(input, axis)
        return np.array([input.sum()]) if axis is None else input.sum(axis=axis)

    @staticmethod
    def backward(ctx, grad_output):
        input, axis = ctx.saved_tensors
        axis = [axis] if type(axis) is int else axis
        shape = [1 if axis is None or i in axis else input.shape[i] for i in range(len(input.shape))]
        return grad_output.reshape(shape) + np.zeros_like(input)


class Max(Function):
    @staticmethod
    def forward(ctx, inp, axis=None):
        axis = [axis] if type(axis) == int else axis
        ret = np.amax(inp, axis=None if axis is None else tuple(axis), keepdims=True)
        ctx.save_for_backward(inp, axis, ret)
        if axis is not None:
            ret = ret.reshape([inp.shape[i] for i in range(len(inp.shape)) if i not in axis])
        return ret

    @staticmethod
    def backward(ctx, grad_output):
        input, axis, ret = ctx.saved_tensors
        shape = [1 if axis is None or i in axis else input.shape[i] for i in range(len(input.shape))]
        ret2 = (input == ret.reshape(shape))
        div = ret2.sum(axis=None if axis is None else tuple(axis), keepdims=True)
        return ret2 * grad_output.reshape(shape) / div


# ************* binary ops *************

def unbroadcast(out, in_sh):
    # adjoint operation to broadcast is sum. Need to sum all axis with 1 = in_sh[i] < out.shape[i]
    sum_axis = tuple([i for i in range(len(in_sh)) if in_sh[i] == 1 and out.shape[i] > 1]) if in_sh != (1,) else None
    return out.sum(axis=sum_axis).reshape(in_sh)


class Add(Function):
    @staticmethod
    def forward(ctx, x, y):
        ctx.save_for_backward(x.shape, y.shape)
        return x + y

    @staticmethod
    def backward(ctx, grad_output):
        shape_x, shape_y = ctx.saved_tensors
        return unbroadcast(grad_output, shape_x), unbroadcast(grad_output, shape_y)


class Sub(Function):
    @staticmethod
    def forward(ctx, x, y):
        ctx.save_for_backward(x.shape, y.shape)
        return x - y

    @staticmethod
    def backward(ctx, grad_output):
        shape_x, shape_y = ctx.saved_tensors
        return unbroadcast(grad_output, shape_x), unbroadcast(-grad_output, shape_y)


class Mul(Function):
    @staticmethod
    def forward(ctx, x, y):
        ctx.save_for_backward(x, y)
        return x * y

    @staticmethod
    def backward(ctx, grad_output):
        x, y = ctx.saved_tensors
        return unbroadcast(y * grad_output, x.shape), unbroadcast(x * grad_output, y.shape)


class Pow(Function):
    @staticmethod
    def forward(ctx, x, y):
        ctx.save_for_backward(x, y)
        return x ** y

    @staticmethod
    def backward(ctx, grad_output):
        x, y = ctx.saved_tensors
        return unbroadcast(y * (x ** (y - 1.0)) * grad_output, x.shape), \
               unbroadcast((x ** y) * np.log(x) * grad_output, y.shape)


# ************* movement ops *************

class Reshape(Function):
    @staticmethod
    def forward(ctx, x, shape):
        ctx.save_for_backward(x.shape)
        return x.reshape(shape)

    @staticmethod
    def backward(ctx, grad_output):
        in_shape, = ctx.saved_tensors
        return grad_output.reshape(in_shape)


class Transpose(Function):
    @staticmethod
    def forward(ctx, x, order):
        ctx.save_for_backward(order)
        return np.transpose(x, order)

    @staticmethod
    def backward(ctx, x):
        return np.transpose(x, np.argsort(ctx.order))


def inner_slice(x, arg):
    padding = [(max(0, -p[0]), max(0, p[1] - x.shape[i])) for i, p in enumerate(arg)]
    x = np.pad(x, padding)
    slicee = [(p[0] + padding[i][0], p[1] + padding[i][0]) for i, p in enumerate(arg)]
    return x[tuple([slice(x[0], x[1], None) for x in slicee])]


class Slice(Function):
    @staticmethod
    def forward(ctx, x, arg=None):
        ctx.save_for_backward(x.shape)
        return inner_slice(x, arg)

    @staticmethod
    def backward(ctx, grad_output):
        shape, = ctx.saved_tensors
        narg = [(0 - p[0], grad_output.shape[i] + (shape[i] - p[1])) for i, p in enumerate(ctx.arg)]
        return inner_slice(grad_output, narg)


# ************* processing ops *************

class Matmul(Function):
    @staticmethod
    def forward(ctx, input, weight):
        ctx.save_for_backward(input, weight)
        return input @ weight

    @staticmethod
    def backward(ctx, grad_output):
        input, weight = ctx.saved_tensors
        grad_input = grad_output @ np.swapaxes(weight, -2, -1)
        grad_weight = np.swapaxes(input, -2, -1) @ grad_output
        return grad_input, grad_weight


class Conv2D(Function):
    @staticmethod
    def forward(ctx, x, w, stride=1, groups=1):
        if type(ctx.stride) == int:
            ctx.stride = (ctx.stride, ctx.stride)
        cout, cin, H, W = w.shape
        ys, xs = ctx.stride
        bs, cin_ = x.shape[0], x.shape[1]
        oy, ox = (x.shape[2] - (H - ys)) // ys, (x.shape[3] - (W - xs)) // xs
        assert cin * ctx.groups == cin_
        assert cout % ctx.groups == 0
        rcout = cout // ctx.groups

        gx = x.reshape(bs, ctx.groups, cin, x.shape[2], x.shape[3])
        tx = np.lib.stride_tricks.as_strided(gx,
                                             shape=(bs, ctx.groups, cin, oy, ox, H, W),
                                             strides=(*gx.strides[0:3], gx.strides[3] * ys, gx.strides[4] * xs,
                                                      *gx.strides[3:5]),
                                             writeable=False,
                                             )
        tw = w.reshape(ctx.groups, rcout, cin, H, W)
        ctx.save_for_backward(tx, tw, x.shape)

        ret = np.zeros((bs, ctx.groups, oy, ox, rcout), dtype=x.dtype)
        for g in range(ctx.groups):
            # ijYXyx,kjyx -> iYXk ->ikYX
            ret[:, g] += np.tensordot(tx[:, g], tw[g], ((1, 4, 5), (1, 2, 3)))
        return np.moveaxis(ret, 4, 2).reshape(bs, cout, oy, ox)

    @staticmethod
    def backward(ctx, grad_output):
        bs, _, oy, ox = grad_output.shape
        tx, tw, x_shape = ctx.saved_tensors
        _, rcout, cin, H, W = tw.shape
        ys, xs = ctx.stride
        OY, OX = x_shape[2:4]

        ggg = grad_output.reshape(bs, ctx.groups, rcout, oy, ox)

        gdw = np.zeros((ctx.groups, rcout, cin, H, W), dtype=tx.dtype)
        for g in range(ctx.groups):
            # 'ikYX,ijYXyx -> kjyx'
            gdw[g] += np.tensordot(ggg[:, g], tx[:, g], ((0, 2, 3), (0, 2, 3)))

        # needs to be optimized
        gdx = np.zeros((bs, ctx.groups, cin, OY, OX), dtype=tx.dtype)
        for k in range(oy * ox):
            Y, X = k // ox, k % ox
            iY, iX = Y * ys, X * xs
            # gdx[:,:,: , iY:iY+H, iX:iX+W] += np.einsum('igk,gkjyx->igjyx', ggg[:,:,:,Y,X], tw)
            for g in range(ctx.groups):
                tg = np.dot(ggg[:, g, :, Y, X].reshape(bs, -1), tw[g].reshape(rcout, -1))
                gdx[:, g, :, iY:iY + H, iX:iX + W] += tg.reshape((bs, cin, H, W))

        return gdx.reshape((bs, ctx.groups * cin, OY, OX)), gdw.reshape((ctx.groups * rcout, cin, H, W))
