"""神经网络层。Module 会自动收集属性里的 Parameter / 子 Module。"""
import numpy as np

from . import functional as F
from .tensor import Tensor, get_default_dtype


class Parameter(Tensor):
    def __init__(self, data):
        super().__init__(np.asarray(data, dtype=get_default_dtype()), requires_grad=True)


class Module:
    training = True

    def forward(self, *args, **kwargs):
        raise NotImplementedError

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

    def named_parameters(self, prefix=""):
        for name, v in vars(self).items():
            items = [(name, v)] if not isinstance(v, (list, tuple)) else [(f"{name}.{i}", x) for i, x in enumerate(v)]
            for n, x in items:
                if isinstance(x, Parameter):
                    yield prefix + n, x
                elif isinstance(x, Module):
                    yield from x.named_parameters(prefix + n + ".")

    def parameters(self):
        return [p for _, p in self.named_parameters()]

    def modules(self):
        yield self
        for v in vars(self).values():
            for x in (v if isinstance(v, (list, tuple)) else [v]):
                if isinstance(x, Module):
                    yield from x.modules()

    def train(self, mode=True):
        for m in self.modules():
            m.training = mode
        return self

    def eval(self):
        return self.train(False)

    def zero_grad(self):
        for p in self.parameters():
            p.grad = None

    def num_parameters(self):
        return sum(p.data.size for p in self.parameters())

    # ---------- 保存 / 加载 ----------
    def state_dict(self):
        return {n: p.data.copy() for n, p in self.named_parameters()}

    def load_state_dict(self, state):
        for n, p in self.named_parameters():
            assert state[n].shape == p.shape, f"{n} 形状不匹配: {state[n].shape} vs {p.shape}"
            p.data = state[n].astype(p.dtype)

    def save(self, path):
        np.savez(path, **self.state_dict())

    def load(self, path):
        with np.load(path) as f:
            self.load_state_dict(dict(f))
        return self


class Linear(Module):
    def __init__(self, in_features, out_features, bias=True):
        bound = 1 / np.sqrt(in_features)
        self.weight = Parameter(np.random.uniform(-bound, bound, (in_features, out_features)))
        self.bias = Parameter(np.zeros(out_features)) if bias else None

    def forward(self, x):
        y = x @ self.weight
        return y + self.bias if self.bias is not None else y


class Embedding(Module):
    def __init__(self, num_embeddings, dim):
        self.weight = Parameter(np.random.randn(num_embeddings, dim) * 0.02)

    def forward(self, idx):
        return self.weight[idx.data if isinstance(idx, Tensor) else np.asarray(idx)]


class LayerNorm(Module):
    def __init__(self, dim, eps=1e-5):
        self.eps = eps
        self.weight = Parameter(np.ones(dim))
        self.bias = Parameter(np.zeros(dim))

    def forward(self, x):
        mu = x.mean(axis=-1, keepdims=True)
        xc = x - mu
        var = (xc * xc).mean(axis=-1, keepdims=True)
        return xc / (var + self.eps).sqrt() * self.weight + self.bias


class Dropout(Module):
    def __init__(self, p=0.1):
        self.p = p

    def forward(self, x):
        return F.dropout(x, self.p, self.training)


class ReLU(Module):
    def forward(self, x):
        return x.relu()


class Tanh(Module):
    def forward(self, x):
        return x.tanh()


class GELU(Module):
    def forward(self, x):
        return F.gelu(x)


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

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

    def __getitem__(self, i):
        return self.layers[i]


class CausalSelfAttention(Module):
    """多头因果自注意力：位置 t 只能看到 <= t 的位置。输入/输出形状 (B, T, C)。"""

    def __init__(self, n_embd, n_head, dropout=0.0):
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.qkv = Linear(n_embd, 3 * n_embd)
        self.proj = Linear(n_embd, n_embd)
        self.attn_drop = Dropout(dropout)
        self.resid_drop = Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        H, D = self.n_head, C // self.n_head
        qkv = self.qkv(x).reshape(B, T, 3, H, D).transpose(2, 0, 3, 1, 4)  # (3, B, H, T, D)
        q, k, v = qkv[0], qkv[1], qkv[2]
        att = (q @ k.transpose(0, 1, 3, 2)) * (1.0 / np.sqrt(D))           # (B, H, T, T)
        att = att.masked_fill(np.triu(np.ones((T, T), dtype=bool), 1), -1e9)
        att = self.attn_drop(F.softmax(att, axis=-1))
        y = (att @ v).transpose(0, 2, 1, 3).reshape(B, T, C)
        return self.resid_drop(self.proj(y))