"""
kiri/nn/pooling.py
"""

import numpy as np
from kiri.backend.detect import BACKEND
from kiri.nn.layers import Module
from kiri.tensor import _CPUTensor

if BACKEND == "mlx":
    import mlx.core as mx


class MaxPool2d(Module):
    """
    2D max pooling.
    Input:  (N, C, H, W)
    Output: (N, C, H//kernel_size, W//kernel_size)
    """
    def __init__(self, kernel_size, stride=None, padding=0):
        super().__init__()
        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
        self.stride      = stride if stride is not None else self.kernel_size
        self.stride      = self.stride if isinstance(self.stride, tuple) else (self.stride, self.stride)
        self.padding     = padding

    def forward(self, x):
        if BACKEND == "mlx":
            # MLX pool expects (N, H, W, C)
            x_nhwc = mx.transpose(x, (0, 2, 3, 1))
            out = mx.max_pool2d(x_nhwc,
                                kernel_size=self.kernel_size,
                                stride=self.stride,
                                padding=self.padding)
            return mx.transpose(out, (0, 3, 1, 2))
        else:
            return _cpu_maxpool(x, self.kernel_size, self.stride, self.padding)

    def __repr__(self):
        return f"MaxPool2d(kernel_size={self.kernel_size}, stride={self.stride})"


class AvgPool2d(Module):
    def __init__(self, kernel_size, stride=None, padding=0):
        super().__init__()
        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
        self.stride      = stride if stride is not None else self.kernel_size
        self.stride      = self.stride if isinstance(self.stride, tuple) else (self.stride, self.stride)
        self.padding     = padding

    def forward(self, x):
        if BACKEND == "mlx":
            x_nhwc = mx.transpose(x, (0, 2, 3, 1))
            out = mx.avg_pool2d(x_nhwc,
                                kernel_size=self.kernel_size,
                                stride=self.stride,
                                padding=self.padding)
            return mx.transpose(out, (0, 3, 1, 2))
        else:
            return _cpu_avgpool(x, self.kernel_size, self.stride, self.padding)

    def __repr__(self):
        return f"AvgPool2d(kernel_size={self.kernel_size}, stride={self.stride})"


class AdaptiveAvgPool2d(Module):
    """Pools to a fixed output size — useful for variable input sizes."""
    def __init__(self, output_size):
        super().__init__()
        self.output_size = output_size if isinstance(output_size, tuple) else (output_size, output_size)

    def forward(self, x):
        if BACKEND == "mlx":
            N, C, H, W = x.shape
            oh, ow = self.output_size
            kH, kW = H // oh, W // ow
            x_nhwc = mx.transpose(x, (0, 2, 3, 1))
            out = mx.avg_pool2d(x_nhwc, kernel_size=(kH, kW), stride=(kH, kW))
            return mx.transpose(out, (0, 3, 1, 2))
        else:
            N, C, H, W = x.data.shape
            oh, ow = self.output_size
            kH, kW = H // oh, W // ow
            return _cpu_avgpool(x, (kH, kW), (kH, kW), 0)


# ─── CPU pool helpers ─────────────────────────────────────────────────────────

def _cpu_maxpool(x_t, kernel_size, stride, padding):
    x = x_t.data
    kH, kW = kernel_size
    sH, sW = stride
    N, C, H, W = x.shape

    if padding > 0:
        x = np.pad(x, ((0,0),(0,0),(padding,padding),(padding,padding)),
                   constant_values=-np.inf)

    H_out = (x.shape[2] - kH) // sH + 1
    W_out = (x.shape[3] - kW) // sW + 1
    out   = np.zeros((N, C, H_out, W_out), dtype=np.float32)
    mask  = np.zeros_like(x)

    for i in range(H_out):
        for j in range(W_out):
            patch = x[:, :, i*sH:i*sH+kH, j*sW:j*sW+kW]
            out[:, :, i, j] = patch.max(axis=(2, 3))
            # record mask for backward
            idx = patch.reshape(N, C, -1).argmax(axis=2)
            ri = idx // kW + i*sH
            ci = idx %  kW + j*sW
            for n in range(N):
                for c in range(C):
                    mask[n, c, ri[n,c], ci[n,c]] += 1

    result = _CPUTensor(out, (x_t,), "maxpool2d")

    def _backward():
        dx = np.zeros_like(x)
        for i in range(H_out):
            for j in range(W_out):
                patch = x[:, :, i*sH:i*sH+kH, j*sW:j*sW+kW]
                idx   = patch.reshape(N, C, -1).argmax(axis=2)
                ri = idx // kW + i*sH
                ci = idx %  kW + j*sW
                for n in range(N):
                    for c in range(C):
                        dx[n, c, ri[n,c], ci[n,c]] += result.grad[n, c, i, j]
        if padding > 0:
            x_t.grad += dx[:, :, padding:-padding, padding:-padding]
        else:
            x_t.grad += dx

    result._backward = _backward
    return result


def _cpu_avgpool(x_t, kernel_size, stride, padding):
    x = x_t.data
    kH, kW = kernel_size
    sH, sW = stride
    N, C, H, W = x.shape

    if padding > 0:
        x = np.pad(x, ((0,0),(0,0),(padding,padding),(padding,padding)))

    H_out = (x.shape[2] - kH) // sH + 1
    W_out = (x.shape[3] - kW) // sW + 1
    out   = np.zeros((N, C, H_out, W_out), dtype=np.float32)

    for i in range(H_out):
        for j in range(W_out):
            out[:, :, i, j] = x[:, :, i*sH:i*sH+kH, j*sW:j*sW+kW].mean(axis=(2,3))

    result = _CPUTensor(out, (x_t,), "avgpool2d")

    def _backward():
        dx = np.zeros_like(x)
        for i in range(H_out):
            for j in range(W_out):
                dx[:, :, i*sH:i*sH+kH, j*sW:j*sW+kW] += \
                    result.grad[:, :, i:i+1, j:j+1] / (kH * kW)
        if padding > 0:
            x_t.grad += dx[:, :, padding:-padding, padding:-padding]
        else:
            x_t.grad += dx

    result._backward = _backward
    return result
