"""
kiri/nn/layers.py
"""

import numpy as np
from kiri.backend.detect import BACKEND
from kiri.tensor import _CPUTensor, Tensor

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


# ─── base ─────────────────────────────────────────────────────────────────────

if BACKEND == "mlx":
    _BaseModule = mlx_nn.Module
else:
    _BaseModule = object


class Module(_BaseModule):
    """Base class for all Kiri layers and models."""

    def __init__(self):
        if BACKEND == "mlx":
            super().__init__()

    def parameters(self):
        if BACKEND == "mlx":
            params = []
            def _collect(obj):
                if isinstance(obj, mx.array):
                    params.append(obj)
                elif isinstance(obj, mlx_nn.Module):
                    for v in vars(obj).values():
                        _collect(v)
                elif isinstance(obj, (list, tuple)):
                    for item in obj:
                        _collect(item)
            for v in vars(self).values():
                _collect(v)
            return params
        else:
            params = []
            for attr in vars(self).values():
                if isinstance(attr, _CPUTensor):
                    params.append(attr)
                elif isinstance(attr, Module):
                    params.extend(attr.parameters())
                elif isinstance(attr, (list, tuple)):
                    for item in attr:
                        if isinstance(item, Module):
                            params.extend(item.parameters())
            return params

    def zero_grad(self):
        if BACKEND != "mlx":
            for p in self.parameters():
                p.grad = np.zeros_like(p.data)

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)


# ─── Linear ───────────────────────────────────────────────────────────────────

class Linear(Module):
    def __init__(self, in_features, out_features, bias=True):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        k = np.sqrt(1.0 / in_features)
        if BACKEND == "mlx":
            self.weight = mx.random.uniform(low=-k, high=k, shape=(out_features, in_features))
            self.bias_val = mx.zeros((out_features,)) if bias else None
        else:
            self.weight = _CPUTensor(
                np.random.uniform(-k, k, (out_features, in_features)).astype(np.float32)
            )
            self.bias_val = _CPUTensor(
                np.zeros(out_features, dtype=np.float32)
            ) if bias else None

    def forward(self, x):
        if BACKEND == "mlx":
            out = x @ self.weight.T
            if self.bias_val is not None:
                out = out + self.bias_val
            return out
        else:
            out = x.matmul(self.weight.T)
            if self.bias_val is not None:
                out = out + self.bias_val
            return out

    def parameters(self):
        p = [self.weight]
        if self.bias_val is not None:
            p.append(self.bias_val)
        return p

    def __repr__(self):
        return f"Linear({self.in_features}, {self.out_features})"


# ─── Conv2d ───────────────────────────────────────────────────────────────────

class Conv2d(Module):
    def __init__(self, in_channels, out_channels, kernel_size,
                 stride=1, padding=0, bias=True):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
        self.stride = stride
        self.padding = padding
        kH, kW = self.kernel_size
        k = np.sqrt(1.0 / (in_channels * kH * kW))
        if BACKEND == "mlx":
            self.weight = mx.random.uniform(low=-k, high=k, shape=(out_channels, in_channels, kH, kW))
            self.bias_val = mx.zeros((out_channels,)) if bias else None
        else:
            self.weight = _CPUTensor(
                np.random.uniform(-k, k, (out_channels, in_channels, kH, kW)).astype(np.float32)
            )
            self.bias_val = _CPUTensor(
                np.zeros(out_channels, dtype=np.float32)
            ) if bias else None

    def forward(self, x):
        if BACKEND == "mlx":
            x_nhwc = mx.transpose(x, (0, 2, 3, 1))
            w_nhwc = mx.transpose(self.weight, (0, 2, 3, 1))
            out = mx.conv2d(x_nhwc, w_nhwc, stride=self.stride, padding=self.padding)
            out = mx.transpose(out, (0, 3, 1, 2))
            if self.bias_val is not None:
                out = out + self.bias_val.reshape(1, -1, 1, 1)
            return out
        else:
            return x.conv2d(self.weight, self.bias_val, self.stride, self.padding)

    def parameters(self):
        p = [self.weight]
        if self.bias_val is not None:
            p.append(self.bias_val)
        return p

    def __repr__(self):
        return f"Conv2d({self.in_channels}, {self.out_channels}, kernel_size={self.kernel_size})"


# ─── BatchNorm1d ──────────────────────────────────────────────────────────────

class BatchNorm1d(Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum
        self._is_training = True
        if BACKEND == "mlx":
            self.weight = mx.ones((num_features,))
            self.bias_val = mx.zeros((num_features,))
        else:
            self.weight = _CPUTensor(np.ones(num_features, dtype=np.float32))
            self.bias_val = _CPUTensor(np.zeros(num_features, dtype=np.float32))
            self.running_mean = np.zeros(num_features, dtype=np.float32)
            self.running_var = np.ones(num_features, dtype=np.float32)

    def forward(self, x):
        if BACKEND == "mlx":
            mean = x.mean(axis=0)
            var = ((x - mean) ** 2).mean(axis=0)
            x_norm = (x - mean) / mx.sqrt(var + self.eps)
            return self.weight * x_norm + self.bias_val
        else:
            if self._is_training:
                mean = x.data.mean(axis=0)
                var = x.data.var(axis=0)
                self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean
                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var
            else:
                mean, var = self.running_mean, self.running_var
            x_norm = _CPUTensor((x.data - mean) / np.sqrt(var + self.eps))
            return self.weight * x_norm + self.bias_val

    def parameters(self):
        return [self.weight, self.bias_val]


# ─── Dropout ──────────────────────────────────────────────────────────────────

class Dropout(Module):
    def __init__(self, p=0.5):
        super().__init__()
        self.p = p
        self._is_training = True

    def forward(self, x):
        if not self._is_training or self.p == 0:
            return x
        if BACKEND == "mlx":
            mask = mx.random.uniform(shape=list(x.shape)) > self.p
            return x * mask / (1 - self.p)
        else:
            mask = (np.random.rand(*x.data.shape) > self.p).astype(np.float32)
            return x * _CPUTensor(mask / (1 - self.p), requires_grad=False)


# ─── Flatten ─────────────────────────────────────────────────────────────────

class Flatten(Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        if BACKEND == "mlx":
            return x.reshape(x.shape[0], -1)
        else:
            return x.reshape(x.shape[0], -1)


# ─── Sequential ──────────────────────────────────────────────────────────────

class Sequential(Module):
    def __init__(self, *layers):
        super().__init__()
        self.layers = list(layers)

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

    def parameters(self):
        params = []
        for layer in self.layers:
            if isinstance(layer, Module):
                params.extend(layer.parameters())
        return params

    def __repr__(self):
        layer_str = "\n  ".join(repr(l) for l in self.layers)
        return f"Sequential(\n  {layer_str}\n)"
