"""
Tiny Llama GLU Lab — Consolidated training library.
One file: model definition, activation registry, stability monitoring,
time tracking, dataset builder, and Trainer factory.
"""

import math
import os
import time
import json
import re
from pathlib import Path
from itertools import chain
from typing import Dict, Callable, Optional, List, Any, Tuple

import torch
import torch.nn as nn
from transformers import (
    LlamaConfig,
    LlamaPreTrainedModel,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    DataCollatorForLanguageModeling,
    AutoTokenizer,
    set_seed,
)
from transformers.models.llama.modeling_llama import (
    LlamaAttention,
    LlamaRMSNorm,
    LlamaRotaryEmbedding,
)
from transformers.modeling_outputs import CausalLMOutputWithPast
from datasets import load_dataset


# =============================================================================
# 1. ACTIVATION REGISTRY
# =============================================================================

class GLUActivationRegistry:
    """Own every gating activation you test. Add new variants in one line."""
    _registry: Dict[str, Callable[[torch.Tensor], torch.Tensor]] = {}

    @classmethod
    def register(cls, name: str, fn: Callable[[torch.Tensor], torch.Tensor]) -> None:
        cls._registry[name] = fn

    @classmethod
    def get(cls, name: str) -> Callable[[torch.Tensor], torch.Tensor]:
        if name not in cls._registry:
            raise KeyError(
                f"Activation '{name}' not found. Available: {list(cls._registry.keys())}"
            )
        return cls._registry[name]


# Built-ins
GLUActivationRegistry.register("silu", nn.functional.silu)
GLUActivationRegistry.register("swish", nn.functional.silu)
GLUActivationRegistry.register("relu", nn.functional.relu)
GLUActivationRegistry.register("gelu", nn.functional.gelu)
GLUActivationRegistry.register("mish", nn.functional.mish)
GLUActivationRegistry.register("sigmoid", torch.sigmoid)
GLUActivationRegistry.register("tanh", torch.tanh)
GLUActivationRegistry.register("elu", nn.functional.elu)
GLUActivationRegistry.register("softplus", nn.functional.softplus)


def bilinear(x: torch.Tensor) -> torch.Tensor:
    """Identity — i.e. no activation at all on the gate branch.

    TinyLlamaMLP's default (non-situglu) forward pass is:
        down_proj(act_fn(gate_proj(x)) * up_proj(x))
    Registering the identity function as act_fn turns that into:
        down_proj(gate_proj(x) * up_proj(x))
    which is a plain bilinear layer: two independent linear projections of
    the input, multiplied elementwise, then projected back down. This is
    the "Bilinear" variant from the GLU Variants paper (Shazeer, 2020) —
    exactly the GLU structure, minus the nonlinearity that normally sits
    on the gate. No new forward-pass code is needed; the existing
    act_fn(gate_proj(x)) * up_proj(x) line does the right thing as soon as
    act_fn is the identity.
    """
    return x


GLUActivationRegistry.register("bilinear", bilinear)


# =============================================================================
# 2. CONFIG
# =============================================================================

class TinyLlamaConfig(LlamaConfig):
    """
    Exact Llama config plus one field: `glu_activation`.
    Enforces pure MHA by requiring num_key_value_heads == num_attention_heads.
    """
    model_type = "tiny_llama"

    def __init__(self, glu_activation: str = "silu", **kwargs):
        super().__init__(**kwargs)
        self.glu_activation = glu_activation
        if self.num_key_value_heads != self.num_attention_heads:
            raise ValueError(
                f"Pure MHA required: num_key_value_heads ({self.num_key_value_heads}) "
                f"must equal num_attention_heads ({self.num_attention_heads})."
            )


# =============================================================================
# 3. MODEL
# =============================================================================

class SiTUGLU(nn.Module):
    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        beta1: float = 4.0,
        beta2: float = 25.0,
    ):
        super().__init__()
        self.beta1 = beta1
        self.beta2 = beta2
        self.W_g = nn.Linear(input_dim, hidden_dim, bias=False)
        self.W_u = nn.Linear(input_dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        g = self.W_g(x)
        gate = self.beta1 * torch.tanh(g / self.beta1) * torch.sigmoid(g)
        up = self.beta2 * torch.tanh(self.W_u(x) / self.beta2)
        return gate * up


class TinyLlamaMLP(nn.Module):
    """GLU MLP with swappable gate activation, including full-block SiTUGLU."""
    def __init__(self, config: TinyLlamaConfig):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size

        if config.glu_activation == "situglu":
            self.act_fn = None
            self.gate_proj = None
            self.up_proj = None
            self.situglu = SiTUGLU(self.hidden_size, self.intermediate_size)
        else:
            self.act_fn = GLUActivationRegistry.get(config.glu_activation)
            self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
            self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
            self.situglu = None

        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.situglu is not None:
            return self.down_proj(self.situglu(x))
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))


class TinyLlamaDecoderLayer(nn.Module):
    def __init__(self, config: TinyLlamaConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx)
        self.mlp = TinyLlamaMLP(config)
        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = LlamaRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs,
    ):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        attn_out = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            position_embeddings=position_embeddings,
        )[0]
        hidden_states = residual + attn_out

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        return (hidden_states,)


def _build_causal_mask(
    attention_mask: Optional[torch.Tensor],
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
) -> torch.Tensor:
    """
    Build the 4D float attention mask that scaled_dot_product_attention expects:
    shape (batch, 1, seq_len, seq_len), dtype matching the model's dtype,
    with 0.0 where attention is allowed and a large negative number where it
    is blocked. This does two jobs at once:
      1. Causal masking — token i can only attend to tokens <= i.
      2. Padding masking — real tokens ignore padding tokens (from attention_mask).
    Without this conversion, a raw 0/1 long tensor gets passed straight into
    SDPA, which raises: "Expected attn_mask dtype to be bool or float ...".
    """
    min_value = torch.finfo(dtype).min

    # causal part: upper triangle (future tokens) is blocked
    causal = torch.full((seq_len, seq_len), fill_value=min_value, dtype=dtype, device=device)
    causal = torch.triu(causal, diagonal=1)
    causal = causal[None, None, :, :]  # (1, 1, seq_len, seq_len)

    if attention_mask is None:
        batch_size = 1
        return causal.expand(batch_size, 1, seq_len, seq_len)

    batch_size = attention_mask.shape[0]
    causal = causal.expand(batch_size, 1, seq_len, seq_len).clone()

    # padding part: positions where attention_mask == 0 are blocked for every query
    padding = attention_mask[:, None, None, :].to(device) == 0  # (batch, 1, 1, seq_len)
    causal = causal.masked_fill(padding, min_value)

    return causal


class TinyLlamaModel(LlamaPreTrainedModel):
    config_class = TinyLlamaConfig

    def __init__(self, config: TinyLlamaConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx
        )
        self.layers = nn.ModuleList(
            [TinyLlamaDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
        )
        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        # RoPE: computed once per forward pass and threaded into every
        # decoder layer's attention. Without this, LlamaAttention receives
        # position_embeddings=None and crashes on `cos, sin = position_embeddings`.
        self.rotary_emb = LlamaRotaryEmbedding(config=config)
        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ):
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if position_ids is None:
            seq_len = inputs_embeds.shape[1]
            position_ids = torch.arange(
                seq_len, device=inputs_embeds.device
            ).unsqueeze(0).expand(inputs_embeds.shape[0], -1)

        hidden_states = inputs_embeds
        # (cos, sin) tuple, shared across all layers this forward pass
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        # Convert the raw 0/1 padding mask into the 4D float causal mask that
        # scaled_dot_product_attention actually expects (bool or float,
        # matching hidden_states dtype). Also enforces causal (left-to-right)
        # attention, which the raw mask alone does not do.
        seq_len = hidden_states.shape[1]
        causal_mask = _build_causal_mask(
            attention_mask, seq_len, hidden_states.dtype, hidden_states.device
        )

        for decoder_layer in self.layers:
            layer_outputs = decoder_layer(
                hidden_states,
                attention_mask=causal_mask,
                position_ids=position_ids,
                position_embeddings=position_embeddings,
            )
            hidden_states = layer_outputs[0]

        hidden_states = self.norm(hidden_states)
        if not return_dict:
            return (hidden_states,)
        return {"last_hidden_state": hidden_states, "hidden_states": None, "attentions": None}


class TinyLlamaForCausalLM(LlamaPreTrainedModel):
    config_class = TinyLlamaConfig
    # Tells save_pretrained() that lm_head.weight and model.embed_tokens.weight
    # are the SAME tensor on purpose (see tie_word_embeddings below), so it
    # doesn't error out when it finds two names pointing at shared memory.
    # NOTE: must be a dict (maps tied key -> source key), not a list —
    # transformers' _get_tied_weight_keys() calls .keys() on this.
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

    def __init__(self, config: TinyLlamaConfig):
        super().__init__(config)
        self.model = TinyLlamaModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        if config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ):
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            return_dict=return_dict,
        )
        hidden_states = outputs["last_hidden_state"] if return_dict else outputs[0]
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(
                shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1)
            )

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=None,
            hidden_states=None,
            attentions=None,
        )

    def prepare_inputs_for_generation(
        self, input_ids, past_key_values=None, attention_mask=None, **kwargs
    ):
        if past_key_values:
            input_ids = input_ids[:, -1:]
        position_ids = kwargs.get("position_ids")
        if attention_mask is not None and position_ids is None:
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -1].unsqueeze(-1)
        return {
            "input_ids": input_ids,
            "position_ids": position_ids,
            "past_key_values": past_key_values,
            "attention_mask": attention_mask,
        }


# =============================================================================
# 4. MONITORING ENGINE
# =============================================================================

class StatsEngine:
    """Compute the unified 6-scalar signature for any tensor."""

    @staticmethod
    def compute(
        tensor: torch.Tensor, user_limit: float, dtype_ratio: float
    ) -> Dict[str, float]:
        with torch.no_grad():
            abs_t = tensor.abs()
            dtype_info = torch.finfo(tensor.dtype)
            dtype_limit = (
                dtype_ratio * dtype_info.max
                if not torch.isinf(torch.tensor(dtype_info.max))
                else float("inf")
            )

            return {
                "norm": tensor.norm(2).item(),
                "mean": tensor.mean().item(),
                "std": tensor.std().item(),
                "max_abs": abs_t.max().item(),
                "frac_near_dtype_limit": (
                    (abs_t > dtype_limit).float().mean().item()
                    if not math.isinf(dtype_limit)
                    else 0.0
                ),
                "frac_near_user_limit": (abs_t > user_limit).float().mean().item(),
            }


class StepAccumulator:
    """
    Stores per-tensor entries, then aggregates to layer-scope or global-scope
    using exact population formulas (no tensor retention).
    """

    def __init__(self):
        # name -> {numel, norm, mean, std, max_abs, frac_near_dtype_limit, frac_near_user_limit}
        self.tensors: Dict[str, Dict[str, float]] = {}

    def add(self, name: str, numel: int, stats: Dict[str, float]):
        # With gradient_accumulation_steps > 1, a single optimizer "step"
        # runs several forward/backward micro-batches, so the same hook
        # (e.g. a given layer's activation, or a given param's grad) fires
        # more than once before clear() is next called. Previously this
        # method did a plain overwrite, silently discarding every
        # micro-batch but the last. Merge instead, using the same exact
        # population formulas _aggregate() uses to combine tensors.
        new_entry = {"numel": numel, **stats}
        existing = self.tensors.get(name)
        self.tensors[name] = (
            new_entry if existing is None else self._merge_entry(existing, new_entry)
        )

    @staticmethod
    def _merge_entry(a: Dict[str, float], b: Dict[str, float]) -> Dict[str, float]:
        total_n = a["numel"] + b["numel"]
        if total_n == 0:
            return a
        norm = math.sqrt(a["norm"] ** 2 + b["norm"] ** 2)
        max_abs = max(a["max_abs"], b["max_abs"])
        mean = (a["mean"] * a["numel"] + b["mean"] * b["numel"]) / total_n
        ex2 = (
            a["numel"] * (a["std"] ** 2 + a["mean"] ** 2)
            + b["numel"] * (b["std"] ** 2 + b["mean"] ** 2)
        ) / total_n
        std = math.sqrt(max(0.0, ex2 - mean ** 2))
        frac_dtype = (
            a["frac_near_dtype_limit"] * a["numel"] + b["frac_near_dtype_limit"] * b["numel"]
        ) / total_n
        frac_user = (
            a["frac_near_user_limit"] * a["numel"] + b["frac_near_user_limit"] * b["numel"]
        ) / total_n
        return {
            "numel": total_n,
            "norm": norm,
            "mean": mean,
            "std": std,
            "max_abs": max_abs,
            "frac_near_dtype_limit": frac_dtype,
            "frac_near_user_limit": frac_user,
        }

    def clear(self):
        self.tensors.clear()

    def _aggregate(self, entries: Dict[str, Dict[str, float]]) -> Dict[str, float]:
        if not entries:
            return {}
        numels = [e["numel"] for e in entries.values()]
        total_n = sum(numels)

        # L2 norm
        norm = math.sqrt(sum(e["norm"] ** 2 for e in entries.values()))
        # Max abs
        max_abs = max(e["max_abs"] for e in entries.values())
        # Weighted mean
        mean = sum(e["mean"] * e["numel"] for e in entries.values()) / total_n
        # Pooled std: sqrt( E[σ² + μ²] - μ_global² )
        ex2 = (
            sum(e["numel"] * (e["std"] ** 2 + e["mean"] ** 2) for e in entries.values())
            / total_n
        )
        std = math.sqrt(max(0.0, ex2 - mean ** 2))
        # Weighted fractions
        frac_dtype = (
            sum(e["frac_near_dtype_limit"] * e["numel"] for e in entries.values())
            / total_n
        )
        frac_user = (
            sum(e["frac_near_user_limit"] * e["numel"] for e in entries.values())
            / total_n
        )

        return {
            "norm": norm,
            "mean": mean,
            "std": std,
            "max_abs": max_abs,
            "frac_near_dtype_limit": frac_dtype,
            "frac_near_user_limit": frac_user,
        }

    def get_global_stats(self) -> Dict[str, float]:
        return self._aggregate(self.tensors)

    def get_layer_stats(self, layer_prefix: str) -> Dict[str, float]:
        entries = {
            k: v for k, v in self.tensors.items() if k.startswith(layer_prefix + ".")
        }
        return self._aggregate(entries)


class HookRegistry:
    """Attach and throttle forward/backward hooks."""

    def __init__(self, model: nn.Module):
        self.model = model
        self.handles: List[torch.utils.hooks.RemovableHandle] = []
        self.active = False

    def attach_forward(
        self,
        module_patterns: List[str],
        accumulator: StepAccumulator,
        user_limit: float,
        dtype_ratio: float,
    ):
        for name, module in self.model.named_modules():
            if any(re.search(p, name) for p in module_patterns):
                h = module.register_forward_hook(
                    self._make_forward_hook(name, accumulator, user_limit, dtype_ratio)
                )
                self.handles.append(h)

    def attach_backward(
        self, accumulator: StepAccumulator, user_limit: float, dtype_ratio: float
    ):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                h = param.register_hook(
                    self._make_backward_hook(
                        f"grad.{name}", accumulator, user_limit, dtype_ratio
                    )
                )
                self.handles.append(h)

    def _make_forward_hook(
        self, module_name: str, accumulator: StepAccumulator, user_limit: float, dtype_ratio: float
    ):
        def hook(module, inp, out):
            if not self.active:
                return

            # Modules can return a tensor, a tuple (take first item), or a
            # dict (e.g. TinyLlamaModel returns {"last_hidden_state": ...}).
            # Pull out the first real tensor we find; skip cleanly if none.
            if isinstance(out, dict):
                out_dict = out
                out = out_dict.get("last_hidden_state")
                if out is None:
                    out = next(
                        (v for v in out_dict.values() if torch.is_tensor(v)), None
                    )
            elif isinstance(out, (tuple, list)):
                out = out[0] if len(out) > 0 else None

            if not torch.is_tensor(out):
                return

            stats = StatsEngine.compute(out.detach(), user_limit, dtype_ratio)
            accumulator.add(f"act.{module_name}", out.numel(), stats)

        return hook

    def _make_backward_hook(
        self, param_name: str, accumulator: StepAccumulator, user_limit: float, dtype_ratio: float
    ):
        def hook(grad):
            if not self.active:
                return
            stats = StatsEngine.compute(grad.detach(), user_limit, dtype_ratio)
            accumulator.add(param_name, grad.numel(), stats)

        return hook

    def set_active(self, active: bool):
        self.active = active

    def clear(self):
        for h in self.handles:
            h.remove()
        self.handles.clear()


class StabilityMonitorCallback(TrainerCallback):
    """
    Full stability instrumentation: grad / param / act statistics
    at global, per-layer, and per-tensor scope.
    """

    def __init__(
        self,
        model: nn.Module,
        monitor_every_n_steps: int = 10,
        module_patterns: Optional[List[str]] = None,
        user_limits: Optional[Dict[str, float]] = None,
        dtype_proximity_ratio: float = 0.9,
        log_scope: Optional[Dict[str, bool]] = None,
        monitor_during_eval: bool = False,
    ):
        self.model = model
        self.monitor_every_n_steps = monitor_every_n_steps
        self.module_patterns = module_patterns or [".*mlp.*", ".*self_attn.*"]
        self.user_limits = user_limits or {"grad": 1.0, "param": 100.0, "act": 50.0}
        self.dtype_ratio = dtype_proximity_ratio
        self.log_scope = log_scope or {
            "global": True,
            "per_layer": True,
            "per_tensor": False,
        }
        self.monitor_during_eval = monitor_during_eval

        self.accumulator = StepAccumulator()
        self.hooks = HookRegistry(model)
        self.hooks.attach_forward(
            self.module_patterns,
            self.accumulator,
            self.user_limits["act"],
            self.dtype_ratio,
        )
        self.hooks.attach_backward(
            self.accumulator, self.user_limits["grad"], self.dtype_ratio
        )

        self.pending_metrics: Optional[Dict[str, float]] = None

    def _should_monitor(self, state) -> bool:
        return state.global_step % self.monitor_every_n_steps == 0

    def on_step_begin(self, args, state, control, **kwargs):
        if self._should_monitor(state):
            self.accumulator.clear()
            self.hooks.set_active(True)

    def on_step_end(self, args, state, control, **kwargs):
        if not self.hooks.active:
            return

        # Parameter stats (post-optimizer step)
        for name, param in self.model.named_parameters():
            stats = StatsEngine.compute(
                param.data, self.user_limits["param"], self.dtype_ratio
            )
            self.accumulator.add(f"param.{name}", param.numel(), stats)

        self.hooks.set_active(False)
        self.pending_metrics = self._build_metrics()

    @staticmethod
    def _kind_of(name: str) -> str:
        """Classify a tensor key by its source: activation, gradient, or parameter."""
        if name.startswith("act."):
            return "act"
        if name.startswith("grad."):
            return "grad"
        if name.startswith("param."):
            return "param"
        return "other"

    @staticmethod
    def _strip_kind(name: str) -> str:
        if name.startswith("act."):
            return name[4:]
        if name.startswith(("grad.", "param.")):
            return name[5:]
        return name

    def _build_metrics(self, scope: str = "train") -> Dict[str, float]:
        metrics: Dict[str, float] = {}

        # --- Global (split by kind: act / grad / param — never pooled together) ---
        if self.log_scope.get("global", True):
            by_kind: Dict[str, Dict[str, Dict[str, float]]] = {}
            for k, v in self.accumulator.tensors.items():
                by_kind.setdefault(self._kind_of(k), {})[k] = v

            for kind, entries in by_kind.items():
                stats = self.accumulator._aggregate(entries)
                for kk, vv in stats.items():
                    metrics[f"{scope}/global/{kind}/{kk}"] = vv

        # --- Per-layer (group by model.layers.{i}, split by kind) ---
        if self.log_scope.get("per_layer", True):
            layer_prefixes = set()
            for name in self.accumulator.tensors:
                clean = self._strip_kind(name)
                parts = clean.split(".")
                for i, p in enumerate(parts):
                    if p == "layers" and i + 1 < len(parts):
                        prefix = ".".join(parts[: i + 2])
                        layer_prefixes.add(prefix)

            for prefix in layer_prefixes:
                by_kind: Dict[str, Dict[str, Dict[str, float]]] = {}
                for k, v in self.accumulator.tensors.items():
                    clean = self._strip_kind(k)
                    if clean.startswith(prefix + ".") or clean == prefix:
                        by_kind.setdefault(self._kind_of(k), {})[k] = v

                safe = prefix.replace(".", "_")
                for kind, entries in by_kind.items():
                    if not entries:
                        continue
                    stats = self.accumulator._aggregate(entries)
                    for kk, vv in stats.items():
                        metrics[f"{scope}/layer_{safe}/{kind}/{kk}"] = vv

        # --- Per-tensor ---
        if self.log_scope.get("per_tensor", False):
            for name, stats in self.accumulator.tensors.items():
                safe = name.replace(".", "_")
                for kk, vv in stats.items():
                    if kk == "numel":
                        continue
                    metrics[f"{scope}/tensor_{safe}/{kk}"] = vv

        return metrics

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is not None and self.pending_metrics is not None:
            logs.update(self.pending_metrics)
            self.pending_metrics = None

    def on_prediction_step(self, args, state, control, **kwargs):
        """Fires once per eval/predict batch. Trainer.evaluate() calls this
        for every batch in the eval loop, then calls self.log(output.metrics)
        (which dispatches on_log to every callback, including the wandb/
        tensorboard reporting callbacks) BEFORE on_evaluate() runs. So to get
        eval-time stats into that same on_log dispatch, we have to build
        pending_metrics here, not in on_evaluate — by the time on_evaluate
        fires, self.log() has already happened and it's too late.
        """
        if not self.monitor_during_eval:
            return
        if not self.hooks.active:
            # First batch of this eval pass: start a fresh accumulation and
            # snapshot parameter stats once (they don't change during eval).
            self.accumulator.clear()
            self.hooks.set_active(True)
            for name, param in self.model.named_parameters():
                stats = StatsEngine.compute(
                    param.data, self.user_limits["param"], self.dtype_ratio
                )
                self.accumulator.add(f"param.{name}", param.numel(), stats)
        self.pending_metrics = self._build_metrics(scope="eval")

    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
        self.hooks.set_active(False)
        self.accumulator.clear()


class TimeTrackerCallback(TrainerCallback):
    """Precise training & eval timing with remaining-time estimates."""

    def __init__(self):
        self.step_start: Optional[float] = None
        self.epoch_start: Optional[float] = None
        self.total_train_time = 0.0
        self.step_times: List[float] = []

    def on_epoch_begin(self, args, state, control, **kwargs):
        self.epoch_start = time.perf_counter()

    def on_step_begin(self, args, state, control, **kwargs):
        self.step_start = time.perf_counter()

    def on_step_end(self, args, state, control, **kwargs):
        if self.step_start is not None:
            dt = time.perf_counter() - self.step_start
            self.step_times.append(dt)
            self.total_train_time += dt
            self.step_start = None

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is None:
            return

        logs["train/total_time_seconds"] = self.total_train_time

        if self.step_times:
            recent = self.step_times[-100:]
            logs["train/time_per_step_avg"] = sum(recent) / len(recent)

        if self.epoch_start is not None:
            logs["train/epoch_time_elapsed"] = time.perf_counter() - self.epoch_start

        if state.max_steps and state.global_step > 0:
            avg = self.total_train_time / state.global_step
            remaining = (state.max_steps - state.global_step) * avg
            logs["train/estimated_remaining_minutes"] = remaining / 60.0

    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
        pass


class MetricsLoggerCallback(TrainerCallback):
    """Persist every logged dict as JSONL in the output dir."""

    def __init__(self, output_dir: str):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.log_file = self.output_dir / "training_log.jsonl"

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is None:
            return
        entry = {
            "step": state.global_step,
            "epoch": state.epoch,
            "timestamp": time.time(),
            **logs,
        }
        with open(self.log_file, "a") as f:
            f.write(json.dumps(entry, default=str) + "\n")


# =============================================================================
# 5. DATA & TRAINER FACTORY
# =============================================================================

def build_dataset(
    tokenizer,
    max_seq_len: int = 512,
    split: str = "train",
    dataset_name: str = "roneneldan/TinyStories",
):
    """Concatenate and chunk TinyStories for causal LM. Fast path with multiprocessing."""
    ds = load_dataset(dataset_name, split=split)

    def tokenize(examples):
        # add_special_tokens=False so we control separators ourselves.
        # Without an explicit boundary token, group_texts() below would
        # concatenate unrelated stories back-to-back with nothing marking
        # where one ends and the next begins, teaching the model spurious
        # cross-document continuations. Append EOS to each example so every
        # packed chunk still carries a clear "new document" signal.
        out = tokenizer(examples["text"], add_special_tokens=False)
        eos_id = tokenizer.eos_token_id
        out["input_ids"] = [ids + [eos_id] for ids in out["input_ids"]]
        if "attention_mask" in out:
            out["attention_mask"] = [mask + [1] for mask in out["attention_mask"]]
        return out

    tokenized = ds.map(
        tokenize,
        batched=True,
        num_proc=4,
        remove_columns=ds.column_names,
        desc=f"Tokenizing {split}",
    )

    def group_texts(examples):
        # chain.from_iterable is O(total) instead of O(n²)
        concatenated = {
            k: list(chain.from_iterable(examples[k])) for k in examples.keys()
        }
        total_length = len(concatenated[list(examples.keys())[0]])
        total_length = (total_length // max_seq_len) * max_seq_len
        result = {
            k: [t[i : i + max_seq_len] for i in range(0, total_length, max_seq_len)]
            for k, t in concatenated.items()
        }
        result["labels"] = result["input_ids"].copy()
        return result

    return tokenized.map(
        group_texts,
        batched=True,
        batch_size=10000,
        num_proc=4,
        desc=f"Chunking {split}",
    )


def create_trainer(
    model,
    tokenizer,
    config: Dict[str, Any],
    train_dataset,
    eval_dataset=None,
):
    """Assemble HF Trainer with all custom callbacks."""
    tc = config.get("training", {})
    mc = config.get("monitor", {})

    args = TrainingArguments(
        output_dir=tc.get("output_dir", "./out"),
        num_train_epochs=tc.get("num_train_epochs", 3),
        per_device_train_batch_size=tc.get("per_device_train_batch_size", 16),
        per_device_eval_batch_size=tc.get("per_device_eval_batch_size", 16),
        gradient_accumulation_steps=tc.get("gradient_accumulation_steps", 4),
        learning_rate=tc.get("learning_rate", 3e-4),
        weight_decay=tc.get("weight_decay", 0.0),
        max_grad_norm=tc.get("max_grad_norm", 1.0),
        optim=tc.get("optim", "adamw_torch"),
        warmup_steps=tc.get("warmup_steps", 0),
        lr_scheduler_type=tc.get("lr_scheduler_type", "cosine"),
        bf16=tc.get("bf16", True),
        logging_steps=tc.get("logging_steps", 10),
        eval_strategy=tc.get("eval_strategy", "steps"),
        eval_steps=tc.get("eval_steps", 500),
        save_strategy=tc.get("save_strategy", "steps"),
        save_steps=tc.get("save_steps", 1000),
        load_best_model_at_end=tc.get("load_best_model_at_end", False),
        report_to=tc.get("report_to", "tensorboard"),
        push_to_hub=tc.get("push_to_hub", False),
        hub_model_id=tc.get("hub_model_id", None),
        # Never commit real tokens to config files. Prefer an explicit value
        # in the config only if someone deliberately put one there; normal
        # case is HF_TOKEN in the environment.
        hub_token=tc.get("hub_token") or os.environ.get("HF_TOKEN"),
        max_steps=tc.get("max_steps", -1),
        seed=tc.get("seed", 42),
        data_seed=tc.get("data_seed", 42),
        remove_unused_columns=False,
    )

    callbacks = [
        TimeTrackerCallback(),
    ]

    if mc.get("enabled", True):
        callbacks.append(
            StabilityMonitorCallback(
                model=model,
                monitor_every_n_steps=mc.get("monitor_every_n_steps", 10),
                module_patterns=mc.get("module_patterns", [".*mlp.*", ".*self_attn.*"]),
                user_limits=mc.get(
                    "user_limits", {"grad": 1.0, "param": 100.0, "act": 50.0}
                ),
                dtype_proximity_ratio=mc.get("dtype_proximity_ratio", 0.9),
                log_scope=mc.get(
                    "log_scope",
                    {"global": True, "per_layer": True, "per_tensor": False},
                ),
                monitor_during_eval=mc.get("monitor_during_eval", False),
            )
        )

    # Must be added LAST: on_log() writes the shared `logs` dict to disk, so
    # every callback that injects keys into that dict (e.g. StabilityMonitor's
    # max_abs/norm/std stats) needs to run BEFORE this one, not after.
    callbacks.append(MetricsLoggerCallback(args.output_dir))

    collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=collator,
        callbacks=callbacks,
    )

    # Trainer.__init__ builds its callback list as:
    #   [DEFAULT_CALLBACKS..., <report_to integrations, e.g. WandbCallback>,
    #    *our callbacks (TimeTracker, StabilityMonitor, MetricsLogger)]
    # All callbacks share the *same* `logs` dict object on on_log(), and are
    # invoked in that list order. That means WandbCallback.on_log() was
    # reading `logs` and shipping it off BEFORE TimeTrackerCallback /
    # StabilityMonitorCallback ever mutated it with their train/global/*,
    # train/layer_*, eval/* keys — so those metrics only ever reached the
    # local training_log.jsonl (written by our MetricsLoggerCallback, which
    # happens to run after within our own sublist) and never wandb.
    # Move every reporting-integration callback to the very end so all of
    # our metric-producing callbacks mutate `logs` first.
    try:
        from transformers.integrations import get_reporting_integration_callbacks

        reporting_types = tuple(get_reporting_integration_callbacks(args.report_to))
    except Exception:
        reporting_types = ()

    if reporting_types:
        handler = trainer.callback_handler
        reporting_cbs = [cb for cb in handler.callbacks if isinstance(cb, reporting_types)]
        other_cbs = [cb for cb in handler.callbacks if not isinstance(cb, reporting_types)]
        handler.callbacks = other_cbs + reporting_cbs

    return trainer