"""Cortex-1 en MLX (Apple Silicon) — equivalente a la implementación NumPy/PyTorch.

Misma arquitectura: GPT pre-norm, RMSNorm, atención causal con QKV fusionado,
SwiGLU, embeddings de posición aprendidos, lm_head atado al embedding.
Paridad verificada: max|Δ logits| < 1e-4 vs el modelo NumPy original.

Uso:  python generate_mlx.py "La capital de Francia es"
"""
from __future__ import annotations

import json
import os

import mlx.core as mx
import mlx.nn as nn

MODEL_FILE = "weights.safetensors"


class Attention(nn.Module):
    def __init__(self, d: int, n_heads: int):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d // n_heads
        self.scale = self.d_head ** -0.5
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.proj = nn.Linear(d, d, bias=False)

    def __call__(self, x: mx.array) -> mx.array:
        B, T, d = x.shape
        qkv = self.qkv(x).reshape(B, T, 3, self.n_heads, self.d_head)
        q, k, v = (qkv[:, :, i].transpose(0, 2, 1, 3) for i in range(3))
        att = (q @ k.transpose(0, 1, 3, 2)) * self.scale          # (B,H,T,T)
        mask = mx.triu(mx.full((T, T), -mx.inf, dtype=att.dtype), k=1)
        att = mx.softmax(att + mask, axis=-1)
        y = (att @ v).transpose(0, 2, 1, 3).reshape(B, T, d)
        return self.proj(y)


class MLP(nn.Module):
    def __init__(self, d: int, d_ff: int):
        super().__init__()
        self.gate_proj = nn.Linear(d, d_ff, bias=False)
        self.up_proj = nn.Linear(d, d_ff, bias=False)
        self.down_proj = nn.Linear(d_ff, d, bias=False)

    def __call__(self, x: mx.array) -> mx.array:
        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))


class Block(nn.Module):
    def __init__(self, d: int, n_heads: int, d_ff: int, eps: float):
        super().__init__()
        self.input_layernorm = nn.RMSNorm(d, eps=eps)
        self.self_attn = Attention(d, n_heads)
        self.post_attention_layernorm = nn.RMSNorm(d, eps=eps)
        self.mlp = MLP(d, d_ff)

    def __call__(self, x):
        x = x + self.self_attn(self.input_layernorm(x))
        return x + self.mlp(self.post_attention_layernorm(x))


class Model(nn.Module):
    def __init__(self, cfg: dict):
        super().__init__()
        d = cfg["hidden_size"]
        self.cfg = cfg
        self.embed_tokens = nn.Embedding(cfg["vocab_size"], d)
        self.embed_positions = nn.Embedding(cfg["max_position_embeddings"], d)
        self.layers = [Block(d, cfg["num_attention_heads"],
                             cfg["intermediate_size"], cfg["rms_norm_eps"])
                       for _ in range(cfg["num_hidden_layers"])]
        self.norm = nn.RMSNorm(d, eps=cfg["rms_norm_eps"])

    def __call__(self, input_ids: mx.array) -> mx.array:
        T = input_ids.shape[1]
        h = self.embed_tokens(input_ids) + self.embed_positions(mx.arange(T))
        for layer in self.layers:
            h = layer(h)
        return self.norm(h)

    def logits(self, input_ids: mx.array) -> mx.array:
        """(B,T) int -> (B,T,V); lm_head atado = embed_tokens."""
        h = self(input_ids)
        w = self.embed_tokens.weight            # (V, d)
        return h @ w.T

    @classmethod
    def load(cls, repo: str) -> "Model":
        with open(os.path.join(repo, "config.json")) as f:
            cfg = json.load(f)
        model = cls(cfg)
        weights = mx.load(os.path.join(repo, MODEL_FILE))
        model.load_weights(list(weights.items()), strict=True)
        return model