"""
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 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 hard_gate_100(x: torch.Tensor) -> torch.Tensor:
    """1 if 0 < x < 100, else 0."""
    return ((x > 0) & (x < 100)).to(x.dtype)


GLUActivationRegistry.register("hard_gate_100", hard_gate_100)


# =============================================================================
# 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]):
        self.tensors[name] = {"numel": numel, **stats}

    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()

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

        # --- Global ---
        if self.log_scope.get("global", True):
            for k, v in self.accumulator.get_global_stats().items():
                metrics[f"train/global/{k}"] = v

        # --- Per-layer (group by model.layers.{i}) ---
        if self.log_scope.get("per_layer", True):
            layer_prefixes = set()
            for name in self.accumulator.tensors:
                clean = name
                if name.startswith("act."):
                    clean = name[4:]
                elif name.startswith(("grad.", "param.")):
                    clean = name[5:]

                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:
                entries = {}
                for k, v in self.accumulator.tensors.items():
                    clean = k
                    if k.startswith("act."):
                        clean = k[4:]
                    elif k.startswith(("grad.", "param.")):
                        clean = k[5:]
                    if clean.startswith(prefix + ".") or clean == prefix:
                        entries[k] = v

                if entries:
                    stats = self.accumulator._aggregate(entries)
                    safe = prefix.replace(".", "_")
                    for kk, vv in stats.items():
                        metrics[f"train/layer_{safe}/{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"train/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_evaluate(self, args, state, control, metrics=None, **kwargs):
        pass


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):
        return tokenizer(examples["text"], add_special_tokens=False)

    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),
        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),
        hub_token=tc.get("hub_token", None),
        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(),
        MetricsLoggerCallback(args.output_dir),
    ]

    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),
            )
        )

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

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