scripts/neuro_trainer.py

#!/usr/bin/env python3
# SPDX-License-Identifier: BUSL-1.1
# Copyright (c) 2024-2026 Lucas Ricardo Mella Chillemi
"""
PAMPAr Neuro-Trainer — Entrenamiento correctivo con losses auxiliares.

Diagnóstico del Brain Scanner reveló 3 problemas en PamparV3:
  1. Tálamo no diferencia — TODOS los tokens van a SINTAXIS (~65-100%)
  2. Early Exit nunca activa — confianza máxima 48% (necesita 90%)
  3. Streams sub-utilizados — la especialización interna no se aprovecha

Este entrenador agrega 3 losses auxiliares para corregir sin cambiar la arquitectura:
  - loss_routing: Supervisión directa con LLAVES → CE contra territorio correcto
  - loss_exit:    Calibra la confianza del Early Exit
  - loss_balance: Previene streams muertos

la loss_routing usa clasificar_token() para determinar el territorio correcto
de cada token, y entrena el routing con CE → gradiente fuerte y direccional
que rompe la simetría del routing uniforme.

Uso:
  python scripts/neuro_trainer.py
  python scripts/neuro_trainer.py --data data/sft_v5.jsonl --pasos 500
  python scripts/neuro_trainer.py --alpha-diff 0.2 --alpha-exit 0.1
"""

from __future__ import annotations

import argparse
import dataclasses
import json
import logging
import math
import random
import sys
import time
from collections import deque
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.utils as nn_utils
import sentencepiece as spm

# Proyecto
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from pampar.coder.v3 import PamparV3, ConfigV3, PRESET_V3
from pampar.coder.v3.talamo import TalamoInicial
from pampar.coder.v3.llaves import clasificar_token
from pampar.coder.v3.zonas import ZONA_TERRITORIO

logger = logging.getLogger("neuro_trainer")
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)

STREAM_NAMES = ["SINTAXIS", "SEMANTICA", "LOGICO", "ESTRUCTURAL"]
_MARCADORES = ["### Solution:", "### Protocolo:"]


# ─────────────────────────────────────────────────────────────────
# Data Loading
# ─────────────────────────────────────────────────────────────────

def _cargar_datos(ruta: Path) -> list[str]:
    """Carga textos desde JSONL."""
    textos: list[str] = []
    with open(ruta, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            texto = obj.get("text", "")
            if texto:
                textos.append(texto)
    logger.info("Cargados %d ejemplos desde %s", len(textos), ruta.name)
    return textos


def _tokenizar_con_mascara(
    textos: list[str],
    tokenizer: spm.SentencePieceProcessor,
    max_len: int,
) -> list[tuple[list[int], list[bool]]]:
    """Tokeniza y crea máscara: loss solo en porción Solution."""
    chunks: list[tuple[list[int], list[bool]]] = []
    for texto in textos:
        ids_full = tokenizer.Encode(texto)
        if len(ids_full) < 8:
            continue
        ids_full = ids_full[: max_len + 1]

        n_prefijo = len(ids_full)
        for marcador in _MARCADORES:
            pos = texto.find(marcador)
            if pos >= 0:
                ids_pre = tokenizer.Encode(texto[: pos + len(marcador)])
                n_prefijo = min(len(ids_pre), len(ids_full))
                break

        mascara = [False] * n_prefijo + [True] * (len(ids_full) - n_prefijo)
        chunks.append((ids_full, mascara))
    return chunks


def _hacer_batch(
    chunks: list[tuple[list[int], list[bool]]],
    indices: list[int],
    device: torch.device,
    max_seq_len: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Batch con padding y máscara de loss."""
    sels = [chunks[i] for i in indices]
    max_len = min(max(len(ids) for ids, _ in sels), max_seq_len + 1)
    padded_ids: list[list[int]] = []
    padded_mask: list[list[bool]] = []
    for ids, mask in sels:
        t = ids[:max_len]
        m = mask[:max_len]
        pad_len = max_len - len(t)
        padded_ids.append(t + [0] * pad_len)
        padded_mask.append(m + [False] * pad_len)
    tokens = torch.tensor(padded_ids, dtype=torch.long, device=device)
    mascara = torch.tensor(padded_mask, dtype=torch.bool, device=device)
    return tokens, mascara


# ─────────────────────────────────────────────────────────────────
# Forward Instrumentado (con gradientes, sin checkpointing)
# ─────────────────────────────────────────────────────────────────

def forward_neuro(
    model: PamparV3,
    input_ids: torch.Tensor,
) -> tuple[torch.Tensor, dict]:
    """
    Forward que captura datos por nivel para las losses auxiliares.

    NO usa gradient checkpointing → guarda todas las activaciones.
    Usar con seq_len cortas y batch=1 para caber en 4GB VRAM.

    Returns:
        logits [B, L, V], dict con datos por nivel
    """
    config = model.config
    n_streams = config.n_streams

    x = model.emb_drop(model.tok_emb(input_ids))
    terr_acts, zona_acts_n0 = model.talamo(x, input_ids)
    streams = [x.clone() for _ in range(n_streams)]

    all_terr_acts: list[torch.Tensor] = [terr_acts]
    all_conf_tensors: list[torch.Tensor] = []
    all_stream_norms: list[torch.Tensor] = []
    all_x_out: list[torch.Tensor] = []

    for i, nivel in enumerate(model.niveles):
        streams, terr_acts, _ = nivel(streams, terr_acts, TalamoInicial.agregar_fn)
        all_terr_acts.append(terr_acts)

        # Recomputar confianza per-token CON gradientes
        # (dentro de nivel.forward() se calcula pero se detacha con .item())
        x_out = sum(
            streams[t] * terr_acts[:, :, t : t + 1] for t in range(n_streams)
        )
        conf = torch.sigmoid(nivel.exit_head(x_out)).squeeze(-1)  # [B, L]
        all_conf_tensors.append(conf)
        all_x_out.append(x_out)

        # Norma L2 promedio por stream (diferenciable)
        norms = torch.stack(
            [streams[t].norm(dim=-1).mean() for t in range(n_streams)]
        )
        all_stream_norms.append(norms)

    # Combinación final + logits
    x_final = model._combinar_streams(streams, terr_acts)
    x_final = model.norm_f(x_final)
    logits = model.lm_head(x_final)

    return logits, {
        "all_terr_acts": all_terr_acts,  # [n_levels+1] × [B, L, 4]
        "all_conf_tensors": all_conf_tensors,  # [n_levels] × [B, L]
        "all_stream_norms": all_stream_norms,  # [n_levels] × [4]
        "all_x_out": all_x_out,  # [n_levels] × [B, L, D]
        "zona_acts_n0": zona_acts_n0,  # [B, L, 52] — zone acts from TalamoInicial
    }


# ─────────────────────────────────────────────────────────────────
# Losses Auxiliares
# ─────────────────────────────────────────────────────────────────

def _build_territory_table(
    tokenizer: spm.SentencePieceProcessor,
) -> torch.Tensor:
    """
    Construye lookup table: token_id → territorio target (0-3).

    Usa clasificar_token() de LLAVES para determinar la zona de cada
    token del vocabulario, y ZONA_TERRITORIO para mapear a territorio.
    Se ejecuta UNA vez al inicio (~48K tokens, <2s).
    """
    vocab_size = tokenizer.GetPieceSize()
    table = torch.zeros(vocab_size, dtype=torch.long)
    for token_id in range(vocab_size):
        piece = tokenizer.IdToPiece(token_id)
        zona, _conf = clasificar_token(piece)
        table[token_id] = ZONA_TERRITORIO[zona].value
    return table


def _build_zone_table(
    tokenizer: spm.SentencePieceProcessor,
) -> torch.Tensor:
    """Lookup table: token_id → zone index (0-51) para loss a nivel de zona."""
    vocab_size = tokenizer.GetPieceSize()
    table = torch.zeros(vocab_size, dtype=torch.long)
    for token_id in range(vocab_size):
        piece = tokenizer.IdToPiece(token_id)
        zona, _conf = clasificar_token(piece)
        table[token_id] = zona.value - 1  # 0-indexed (0-51)
    return table


def calcular_loss_talamo(
    all_terr_acts: list[torch.Tensor],
    zona_acts_n0: torch.Tensor,
    input_ids: torch.Tensor,
    territory_table: torch.Tensor,
    zone_table: torch.Tensor,
) -> torch.Tensor:
    """
    Loss enfocada en TalamoInicial: CE a nivel de zona + CE a nivel de territorio.

    La zona-level CE bypasses la dilución del MEAN en agregar_zonas_a_territorios,
    dando gradiente directo a attn_proj y context_conv para activar las zonas
    correctas. Sin esto, el gradiente a través de MEAN(15 zonas) es ~0.013x.
    """
    device = input_ids.device

    # CE sobre territorio en N0 (4-class)
    targets_terr = territory_table.to(device)[input_ids]
    terr_acts_n0 = all_terr_acts[0]
    loss_terr = F.cross_entropy(
        terr_acts_n0.reshape(-1, 4), targets_terr.reshape(-1)
    )

    # CE sobre zonas en N0 (52-class) — bypasses MEAN dilution
    targets_zona = zone_table.to(device)[input_ids]
    loss_zona = F.cross_entropy(
        zona_acts_n0.reshape(-1, 52), targets_zona.reshape(-1)
    )

    return loss_terr + 2.0 * loss_zona


def calcular_loss_routing(
    all_terr_acts: list[torch.Tensor],
    input_ids: torch.Tensor,
    territory_table: torch.Tensor,
) -> torch.Tensor:
    """
    Supervised routing loss: CE contra el territorio correcto de LLAVES.

    Cada token tiene un territorio correcto determinado por clasificar_token().
    Usa class weights inversamente proporcionales a la frecuencia para compensar
    el desbalance masivo SEMA=47751 vs SINT=169 vs LOGI=72 vs ESTR=8.
    Sin esto, el gradiente para SINT/LOGI tokens es demasiado débil.
    """
    device = input_ids.device
    # SINT=169, SEMA=47751, LOGI=72, ESTR=8 → weights= sqrt(total/count), capped at 10
    weights = torch.tensor(
        [min(10.0, (48000 / 169) ** 0.5),   # SINT:  ~16.9 → 10
         1.0,                                 # SEMA:  ~1.0
         min(10.0, (48000 / 72) ** 0.5),     # LOGI:  ~25.8 → 10
         min(10.0, (48000 / 8) ** 0.5)],     # ESTR:  ~77.5 → 10
        device=device,
        dtype=torch.float32,
    )

    targets = territory_table.to(device)[input_ids]  # [B, L]

    total = torch.tensor(0.0, device=device)
    for ta in all_terr_acts:
        L_ta = ta.shape[1]
        t = targets[:, :L_ta]
        total = total + F.cross_entropy(
            ta.reshape(-1, 4), t.reshape(-1), weight=weights
        )
    return total / len(all_terr_acts)


def calcular_loss_exit(
    all_conf_tensors: list[torch.Tensor],
    all_intermediate_correct: list[torch.Tensor],
    loss_mask: torch.Tensor,
) -> torch.Tensor:
    """
    Loss de calibración de Early Exit en TODOS los niveles.

    Para cada nivel: BCE entre su confianza y si predice correctamente
    (usando el LM head compartido en representaciones intermedias).
    Niveles profundos tienen más peso (predicen mejor).
    """
    device = all_conf_tensors[0].device
    total = torch.tensor(0.0, device=device)
    n_levels = len(all_conf_tensors)
    valid = loss_mask.float()
    n_valid = valid.sum().clamp(min=1)

    # — Calibración en CADA nivel (peso creciente con profundidad) —
    for i in range(n_levels):
        conf = all_conf_tensors[i]  # [B, L]
        correct = all_intermediate_correct[i]  # [B, L]
        weight = (i + 1) / n_levels  # 0.2, 0.4, 0.6, 0.8, 1.0
        total = total + weight * F.binary_cross_entropy(
            conf.clamp(1e-7, 1 - 1e-7) * valid,
            correct * valid,
            reduction="sum",
        ) / n_valid

    # — Monotonía: confianza debe subir con profundidad —
    for i in range(1, n_levels):
        violation = all_conf_tensors[i - 1] - all_conf_tensors[i]
        total = total + F.relu(violation).mean()

    return total


def calcular_loss_balance(all_stream_norms: list[torch.Tensor]) -> torch.Tensor:
    """
    Loss de utilización de streams.

    Penaliza el stream con menor norma → previene streams muertos.
    Todos los streams deben contribuir, no solo SINTAXIS.
    """
    total = torch.tensor(0.0, device=all_stream_norms[0].device)
    for norms in all_stream_norms:
        min_norm = norms.min()
        total = total - torch.log(min_norm.clamp(min=1e-8))
    return total / len(all_stream_norms)


def calcular_loss_norm(
    all_stream_norms: list[torch.Tensor],
    max_norms: list[float],
) -> torch.Tensor:
    """
    Soft norm penalty: penaliza streams cuya norma supera el max_norm por nivel.

    Enseña al modelo a mantener normas controladas SIN hard clamp.
    max_norms[i] = límite suave del nivel i (50, 100, 200, 400, 800).
    """
    total = torch.tensor(0.0, device=all_stream_norms[0].device)
    for i, norms in enumerate(all_stream_norms):
        limit = max_norms[i] if i < len(max_norms) else max_norms[-1]
        excess = F.relu(norms - limit)  # [n_streams]
        total = total + (excess ** 2).mean()
    return total / len(all_stream_norms)


# ─────────────────────────────────────────────────────────────────
# Training Step
# ─────────────────────────────────────────────────────────────────

def paso_neuro(
    modelo: PamparV3,
    optimizer: torch.optim.Optimizer,
    tokens: torch.Tensor,
    mascara: torch.Tensor,
    max_grad_norm: float,
    alpha_diff: float,
    alpha_exit: float,
    alpha_balance: float,
    alpha_norm: float,
    territory_table: torch.Tensor,
    warmup_factor: float = 1.0,
) -> dict[str, float]:
    """Un paso de entrenamiento con CE + 4 losses auxiliares."""
    modelo.train()
    optimizer.zero_grad(set_to_none=True)

    input_ids = tokens[:, :-1]
    targets = tokens[:, 1:]
    loss_mask = mascara[:, 1:]

    # Forward instrumentado (sin checkpointing)
    logits, neuro_info = forward_neuro(modelo, input_ids)

    B, T, V = logits.shape

    # CE Loss (solo en solución)
    targets_masked = targets.masked_fill(~loss_mask, -100)
    loss_ce = F.cross_entropy(
        logits.reshape(B * T, V),
        targets_masked.reshape(B * T),
        ignore_index=-100,
    )

    # Losses auxiliares (rampean con warmup_factor)
    loss_diff = calcular_loss_routing(
        neuro_info["all_terr_acts"], input_ids, territory_table
    )

    # Intermediate correctness para calibrar exit en TODOS los niveles
    with torch.no_grad():
        all_intermediate_correct: list[torch.Tensor] = []
        for x_out in neuro_info["all_x_out"]:
            inter_logits = modelo.lm_head(modelo.norm_f(x_out))
            inter_pred = inter_logits.argmax(dim=-1)
            correct = (inter_pred == targets_masked).float()
            all_intermediate_correct.append(correct)

    loss_exit = calcular_loss_exit(
        neuro_info["all_conf_tensors"], all_intermediate_correct, loss_mask
    )
    loss_balance = calcular_loss_balance(neuro_info["all_stream_norms"])

    # Norm penalty: enseña al modelo a mantener normas controladas
    max_norms = [50.0 * (2.0 ** i) for i in range(len(neuro_info["all_stream_norms"]))]
    loss_norm = calcular_loss_norm(neuro_info["all_stream_norms"], max_norms)

    wf = warmup_factor
    loss = (
        loss_ce
        + wf * alpha_diff * loss_diff
        + wf * alpha_exit * loss_exit
        + wf * alpha_balance * loss_balance
        + wf * alpha_norm * loss_norm
    )

    if loss.isnan() or loss.isinf():
        logger.warning("Loss inestable — skipping step")
        return {"ce": 0.0, "diff": 0.0, "exit": 0.0, "balance": 0.0, "norm": 0.0, "total": 0.0}

    loss.backward()
    nn_utils.clip_grad_norm_(modelo.parameters(), max_grad_norm)
    optimizer.step()

    return {
        "ce": float(loss_ce.detach()),
        "diff": float(loss_diff.detach()),
        "exit": float(loss_exit.detach()),
        "balance": float(loss_balance.detach()),
        "norm": float(loss_norm.detach()),
        "total": float(loss.detach()),
    }


# ─────────────────────────────────────────────────────────────────
# Diagnóstico Periódico
# ─────────────────────────────────────────────────────────────────

@torch.no_grad()
def diagnostico(
    modelo: PamparV3,
    tokenizer: spm.SentencePieceProcessor,
    device: torch.device,
    territory_table: torch.Tensor | None = None,
) -> str:
    """Mini-scan con frase de test para ver si las losses están funcionando."""
    modelo.eval()

    test_code = "def fibonacci(n):"
    ids = tokenizer.Encode(test_code)
    input_ids = torch.tensor([ids], device=device)

    _, info = forward_neuro(modelo, input_ids)

    lines = ["  ── Diagnóstico ──"]

    # Routing diversity: std de terr_acts por nivel
    for lvl, ta in enumerate(info["all_terr_acts"]):
        std_mean = torch.std(ta, dim=-1).mean().item()
        lines.append(f"    Nivel {lvl} routing std: {std_mean:.4f}")

    # Confianza por nivel
    confs = [c.mean().item() for c in info["all_conf_tensors"]]
    conf_str = " → ".join(f"{c:.3f}" for c in confs)
    lines.append(f"    Confianza: {conf_str}")

    # Stream balance: normas del último nivel
    last_norms = info["all_stream_norms"][-1]
    norm_str = " | ".join(
        f"{STREAM_NAMES[t][:4]}={last_norms[t]:.1f}" for t in range(4)
    )
    lines.append(f"    Streams: {norm_str}")

    # Routing dominante (actual)
    final_ta = info["all_terr_acts"][-1]  # [1, L, 4]
    dom_counts = [0, 0, 0, 0]
    for i in range(final_ta.shape[1]):
        dom = final_ta[0, i].argmax().item()
        dom_counts[dom] += 1
    dom_str = " | ".join(
        f"{STREAM_NAMES[t][:4]}={dom_counts[t]}" for t in range(4)
    )
    lines.append(f"    Actual:   {dom_str}")

    # Routing esperado (LLAVES target)
    if territory_table is not None:
        expected = territory_table[input_ids[0].cpu()]
        exp_counts = [0, 0, 0, 0]
        for e in expected:
            exp_counts[e.item()] += 1
        exp_str = " | ".join(
            f"{STREAM_NAMES[t][:4]}={exp_counts[t]}" for t in range(4)
        )
        lines.append(f"    Esperado: {exp_str}")

    modelo.train()
    return "\n".join(lines)


# ─────────────────────────────────────────────────────────────────
# Utils
# ─────────────────────────────────────────────────────────────────

def _cosine_lr(
    paso: int, warmup: int, total: int, lr_max: float, lr_min: float
) -> float:
    if paso < warmup:
        return lr_max * (paso + 1) / warmup
    progreso = (paso - warmup) / max(1, total - warmup)
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * progreso))


def _guardar(
    ruta: Path, modelo: PamparV3, optimizer: torch.optim.Optimizer, paso: int
) -> None:
    ruta.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "modelo": modelo.state_dict(),
            "optimizer": optimizer.state_dict(),
            "paso_global": paso,
            "config": dataclasses.asdict(modelo.config),
            "tipo": "neuro_trainer",
        },
        ruta,
    )
    logger.info("Checkpoint guardado → %s (paso %d)", ruta.name, paso)


# ─────────────────────────────────────────────────────────────────
# CLI + main
# ─────────────────────────────────────────────────────────────────

def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="PAMPAr Neuro-Trainer: corrige routing, exit y balance",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument(
        "--checkpoint-in",
        type=Path,
        default=ROOT / "checkpoints" / "v3_sft_v8.pt",
    )
    p.add_argument(
        "--checkpoint-out",
        type=Path,
        default=ROOT / "checkpoints" / "v3_neuro_v1.pt",
    )
    p.add_argument(
        "--tokenizer",
        type=Path,
        default=ROOT / "data" / "tokenizer" / "pampar_48k.model",
    )
    p.add_argument("--data", type=Path, default=ROOT / "data" / "sft_v5.jsonl")

    p.add_argument("--lr", type=float, default=3e-7,
                   help="LR para params principales (muy bajo, proteger CE)")
    p.add_argument("--lr-min", type=float, default=3e-8)
    p.add_argument("--lr-routing", type=float, default=1e-3,
                   help="LR para params de routing (talamo, lateral.scale)")
    p.add_argument("--lr-routing-min", type=float, default=1e-4)
    p.add_argument("--warmup", type=int, default=30)
    p.add_argument("--max-pasos", type=int, default=1000)
    p.add_argument("--epochs", type=int, default=20)
    p.add_argument("--batch-size", type=int, default=1,
                   help="Batch=1 (4GB VRAM, sin checkpointing)")
    p.add_argument("--seq-len", type=int, default=256,
                   help="Secuencia corta para caber sin checkpointing")
    p.add_argument("--max-grad-norm", type=float, default=1.0)
    p.add_argument("--guardar-cada", type=int, default=200)

    # Pesos de losses auxiliares
    p.add_argument("--alpha-diff", type=float, default=0.5,
                   help="Peso loss_routing (CE supervisado contra LLAVES)")
    p.add_argument("--alpha-exit", type=float, default=0.0,
                   help="Peso loss_exit (desactivado hasta resolver routing)")
    p.add_argument("--alpha-balance", type=float, default=0.05,
                   help="Peso loss_balance (utilizaci\u00f3n streams)")
    p.add_argument("--alpha-norm", type=float, default=0.01,
                   help="Peso loss_norm (penaliza normas excesivas por nivel)")
    p.add_argument("--aux-warmup", type=int, default=50,
                   help="Pasos para rampear aux losses de 0→1")

    # Modo Focus-Talamo (Round 6)
    p.add_argument("--focus-talamo", action="store_true",
                   help="Solo entrena TalamoInicial (attn_proj, conv, gate)")
    p.add_argument("--lr-talamo", type=float, default=5e-3,
                   help="LR para TalamoInicial en modo focus-talamo")

    # Norm clamping en training
    p.add_argument("--train-norm-clamp", action="store_true",
                   help="Activa norm clamping durante training (regulariza normas)")

    p.add_argument("--device", type=str, default="auto")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--diag-cada", type=int, default=50,
                   help="Cada cuántos pasos mostrar diagnóstico")

    return p.parse_args()


def main() -> None:
    args = _parse_args()
    random.seed(args.seed)
    torch.manual_seed(args.seed)

    device = torch.device(
        "cuda"
        if args.device == "auto" and torch.cuda.is_available()
        else args.device
        if args.device != "auto"
        else "cpu"
    )
    logger.info("Device: %s", device)
    if device.type == "cuda":
        torch.cuda.manual_seed(args.seed)
        logger.info(
            "GPU: %s (%.1f GiB)",
            torch.cuda.get_device_name(0),
            torch.cuda.get_device_properties(0).total_memory / 1e9,
        )

    # Tokenizer
    tok = spm.SentencePieceProcessor()
    tok.Load(str(args.tokenizer))
    logger.info("Tokenizer vocab=%d", tok.GetPieceSize())

    # Modelo
    if not args.checkpoint_in.exists():
        logger.error("Checkpoint no encontrado: %s", args.checkpoint_in)
        sys.exit(1)

    payload = torch.load(
        args.checkpoint_in, map_location=device, weights_only=False
    )
    config = ConfigV3(**payload["config"]) if "config" in payload else PRESET_V3
    modelo = PamparV3(config).to(device)
    modelo.load_state_dict(payload["modelo"])
    modelo.registrar_tokenizer(tok)
    logger.info(
        "Cargado '%s' (tipo: %s)", args.checkpoint_in.name, payload.get("tipo", "?")
    )
    del payload

    n_params = sum(p.numel() for p in modelo.parameters() if p.requires_grad)
    logger.info("PamparV3 %.1fM params", n_params / 1e6)

    # Norm clamping durante training (regulariza normas por nivel)
    if args.train_norm_clamp:
        modelo.set_train_norm_clamp(True)
        logger.info("Norm clamping ACTIVADO durante training")

    # Datos
    if not args.data.exists():
        logger.error("Dataset no encontrado: %s", args.data)
        sys.exit(1)

    textos = _cargar_datos(args.data)
    chunks = _tokenizar_con_mascara(textos, tok, args.seq_len)
    logger.info("Chunks tokenizados: %d (seq_len=%d)", len(chunks), args.seq_len)
    del textos

    pct_solution = sum(sum(m) for _, m in chunks) / max(
        1, sum(len(ids) for ids, _ in chunks)
    )
    logger.info("Porcentaje tokens en zona Solution: %.1f%%", pct_solution * 100)

    # Territory table para supervised routing loss
    territory_table = _build_territory_table(tok)
    counts = [(territory_table == t).sum().item() for t in range(4)]
    logger.info(
        "Territory table: SINT=%d SEMA=%d LOGI=%d ESTR=%d (total=%d)",
        *counts, sum(counts),
    )

    # Diagnóstico inicial
    logger.info("─── DIAGNÓSTICO INICIAL ───")
    print(diagnostico(modelo, tok, device, territory_table))

    # ── MODO FOCUS-TALAMO (Round 6) ──
    if args.focus_talamo:
        logger.info("═══ MODO FOCUS-TALAMO: entrenando solo TalamoInicial ═══")
        zone_table = _build_zone_table(tok)

        # Congelar TODO excepto TalamoInicial
        for name, p in modelo.named_parameters():
            if not name.startswith("talamo."):
                p.requires_grad_(False)

        n_trainable = sum(p.numel() for p in modelo.parameters() if p.requires_grad)
        logger.info("Params entrenables: %.1fK (TalamoInicial)", n_trainable / 1e3)

        talamo_params = [p for p in modelo.talamo.parameters() if p.requires_grad]
        opt = torch.optim.AdamW(
            talamo_params, lr=args.lr_talamo,
            betas=(0.9, 0.95), weight_decay=0.01,
        )

        t0 = time.time()
        losses_buf: deque[float] = deque(maxlen=100)
        paso = 0

        try:
            for epoch in range(args.epochs):
                idx_all = list(range(len(chunks)))
                random.shuffle(idx_all)
                logger.info("── Epoch %d/%d ──", epoch + 1, args.epochs)

                for i in range(0, len(idx_all) - args.batch_size + 1, args.batch_size):
                    batch_idx = idx_all[i : i + args.batch_size]
                    tokens, _ = _hacer_batch(chunks, batch_idx, device, args.seq_len)

                    opt.zero_grad(set_to_none=True)
                    input_ids = tokens[:, :-1]
                    _, info = forward_neuro(modelo, input_ids)

                    loss = calcular_loss_talamo(
                        info["all_terr_acts"], info["zona_acts_n0"],
                        input_ids, territory_table, zone_table,
                    )

                    if loss.isnan() or loss.isinf():
                        continue

                    loss.backward()
                    nn_utils.clip_grad_norm_(modelo.parameters(), args.max_grad_norm)
                    opt.step()

                    losses_buf.append(loss.item())
                    paso += 1

                    if paso % 10 == 0:
                        avg = sum(losses_buf) / len(losses_buf)
                        elapsed = time.time() - t0
                        logger.info(
                            "paso %4d/%d | loss=%.4f  avg=%.4f  (%.1f p/s)",
                            paso, args.max_pasos, loss.item(), avg,
                            paso / elapsed,
                        )

                    if paso % args.diag_cada == 0:
                        print(diagnostico(modelo, tok, device, territory_table))

                    if paso % args.guardar_cada == 0:
                        _guardar(args.checkpoint_out, modelo, opt, paso)

                    if paso >= args.max_pasos:
                        break

                if paso >= args.max_pasos:
                    break

        except KeyboardInterrupt:
            logger.info("Interrumpido — guardando...")

        _guardar(args.checkpoint_out, modelo, opt, paso)
        logger.info("─── DIAGNÓSTICO FINAL (TALAMO) ───")
        print(diagnostico(modelo, tok, device, territory_table))

        elapsed = time.time() - t0
        avg = sum(losses_buf) / max(1, len(losses_buf))
        print(f"\n── Focus-Talamo Completado ──")
        print(f"  Pasos: {paso}")
        print(f"  Tiempo: {int(elapsed // 60)}m{int(elapsed % 60):02d}s")
        print(f"  Loss final (avg): {avg:.4f}")
        print(f"  Checkpoint: {args.checkpoint_out}")
        print(f"\n  Verificar con Brain Scanner:")
        print(
            f'    python -X utf8 scripts/brain_scanner.py --suite --device cuda'
            f' --checkpoint {args.checkpoint_out}'
        )
        return

    # Separar parámetros: routing (LR alto) vs resto (LR bajo)
    routing_names = {"talamo", "talamo_nivel", "lateral"}
    routing_params: list[torch.nn.Parameter] = []
    main_params: list[torch.nn.Parameter] = []
    for name, p in modelo.named_parameters():
        parts = name.split(".")
        if any(rn in parts for rn in routing_names):
            routing_params.append(p)
        else:
            main_params.append(p)

    n_routing = sum(p.numel() for p in routing_params)
    n_main = sum(p.numel() for p in main_params)
    logger.info(
        "Param groups: routing=%.1fK (lr=%.1e) | main=%.1fM (lr=%.1e)",
        n_routing / 1e3, args.lr_routing,
        n_main / 1e6, args.lr,
    )

    optimizer = torch.optim.AdamW(
        [
            {"params": main_params, "lr": args.lr},
            {"params": routing_params, "lr": args.lr_routing},
        ],
        betas=(0.9, 0.95),
        weight_decay=0.01,
        eps=1e-8,
    )

    pasos_por_epoch = max(1, len(chunks) // args.batch_size)
    total_pasos = min(args.max_pasos, args.epochs * pasos_por_epoch)
    logger.info(
        "Neuro-Training: %d pasos | %d chunks | alphas: diff=%.3f exit=%.3f balance=%.3f norm=%.3f",
        total_pasos,
        len(chunks),
        args.alpha_diff,
        args.alpha_exit,
        args.alpha_balance,
        args.alpha_norm,
    )

    paso = 0
    t0 = time.time()
    losses_ce: deque[float] = deque(maxlen=100)
    losses_total: deque[float] = deque(maxlen=100)

    try:
        for epoch in range(args.epochs):
            idx = list(range(len(chunks)))
            random.shuffle(idx)
            logger.info("── Epoch %d/%d ──", epoch + 1, args.epochs)

            for i in range(0, len(idx) - args.batch_size + 1, args.batch_size):
                batch_idx = idx[i : i + args.batch_size]
                tokens, mascara = _hacer_batch(
                    chunks, batch_idx, device, args.seq_len
                )

                lr_main = _cosine_lr(paso, args.warmup, total_pasos, args.lr, args.lr_min)
                lr_rout = _cosine_lr(paso, args.warmup, total_pasos, args.lr_routing, args.lr_routing_min)
                optimizer.param_groups[0]["lr"] = lr_main
                optimizer.param_groups[1]["lr"] = lr_rout

                wf = min(1.0, paso / max(1, args.aux_warmup))

                loss_dict = paso_neuro(
                    modelo,
                    optimizer,
                    tokens,
                    mascara,
                    args.max_grad_norm,
                    args.alpha_diff,
                    args.alpha_exit,
                    args.alpha_balance,
                    args.alpha_norm,
                    territory_table,
                    warmup_factor=wf,
                )

                if loss_dict["ce"] > 0:
                    losses_ce.append(loss_dict["ce"])
                    losses_total.append(loss_dict["total"])

                paso += 1

                if paso % 10 == 0:
                    avg_ce = sum(losses_ce) / max(1, len(losses_ce))
                    avg_total = sum(losses_total) / max(1, len(losses_total))
                    elapsed = time.time() - t0
                    logger.info(
                        "paso %4d/%d | CE=%.3f  diff=%.3f  exit=%.3f  bal=%.3f  "
                        "norm=%.3f  total=%.3f  lr_r=%.1e  (%.1f p/s)",
                        paso,
                        total_pasos,
                        loss_dict["ce"],
                        loss_dict["diff"],
                        loss_dict["exit"],
                        loss_dict["balance"],
                        loss_dict["norm"],
                        avg_total,
                        lr_rout,
                        paso / elapsed,
                    )

                if paso % args.diag_cada == 0:
                    print(diagnostico(modelo, tok, device, territory_table))

                if paso % args.guardar_cada == 0:
                    _guardar(args.checkpoint_out, modelo, optimizer, paso)

                if paso >= args.max_pasos:
                    break

            if paso >= args.max_pasos:
                break

    except KeyboardInterrupt:
        logger.info("Interrumpido — guardando...")

    _guardar(args.checkpoint_out, modelo, optimizer, paso)

    # Diagnóstico final
    logger.info("─── DIAGNÓSTICO FINAL ───")
    print(diagnostico(modelo, tok, device, territory_table))

    elapsed = time.time() - t0
    avg_ce = sum(losses_ce) / max(1, len(losses_ce))
    print(f"\n── Neuro-Training Completado ──")
    print(f"  Pasos: {paso}")
    print(f"  Tiempo: {int(elapsed // 3600)}h{int((elapsed % 3600) // 60):02d}m")
    print(f"  Loss CE final (avg100): {avg_ce:.3f}")
    print(f"  PPL final: {math.exp(min(avg_ce, 10)):.1f}")
    print(f"  Checkpoint: {args.checkpoint_out}")
    print(f"\n  Verificar con Brain Scanner:")
    print(
        f'    python -X utf8 scripts/brain_scanner.py --code "def fibonacci(n):" --device cuda'
    )


if __name__ == "__main__":
    main()