"""Gemma 4 text decoder with LoRA and one linear head per field (needs torch; runs on the pod and on node12).

Classifier-plan §1, §4. The checkpoints (``google/gemma-4-12B`` and ``-it``) are
``Gemma4UnifiedForConditionalGeneration`` (model_type ``gemma4_unified``): text decoder plus light image/audio
projections. Only the text decoder is loaded, as ``Gemma4UnifiedForCausalLM`` with the text config and a
``key_mapping`` that renames the checkpoint's ``model.language_model.*`` to ``model.*`` (without it every weight is
reported missing and the model would train from random weights; the yt sibling found this on 2026-10-18). Image and
audio weights are ignored, the tied LM head is dropped: ``HeadsModel`` keeps only ``lm.model`` (the final-norm
hidden states) and one ``Linear(hidden -> K_f)`` per taxonomy field.

Pooling (the §4.4 screen): ``last`` = the hidden state of each row's last real token (the ``### LABELS`` suffix);
``mean`` = the mean of final-norm states over the row's real tokens. Rows are right-padded with an attention mask.
Gemma needs its BOS token: inputs are built by ``training.formats.encode`` with ``bos_id(tok)``.

The loss (``step_loss``) is normalised per optimiser step: per field, the summed cross-entropy over unmasked rows
divided by that step's unmasked count, fields weighted by ``Taxonomy.weights()`` (direction 2.0). Splitting a step
into micro-batches therefore leaves its gradient unchanged. Optional per-option class weights
(``Taxonomy.class_weight_vectors()``; nsb-v2: pnl breakeven 3.0) scale each row's cross-entropy by the weight of its
gold option; the divisor stays the unweighted count, so the split invariance holds.
"""
from __future__ import annotations

import json
import os
import random
import shutil
from pathlib import Path

import torch
import torch.nn.functional as F
from torch import nn

LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
# model_type -> text-only class and checkpoint key renames.
BACKBONES = {
    "gemma4_unified": dict(name="gemma4", cls="Gemma4UnifiedForCausalLM",
                           key_mapping={r"^model\.language_model\.": "model."}, bos=True),
    "gemma4_unified_text": dict(name="gemma4", cls="Gemma4UnifiedForCausalLM", key_mapping=None, bos=True),
}
HUB_IDS = {"base": "google/gemma-4-12B", "it": "google/gemma-4-12B-it"}
POOLS = ("last", "mean")


def backbone(path) -> dict:
    from transformers import AutoConfig

    model_type = AutoConfig.from_pretrained(path).model_type
    if model_type not in BACKBONES:
        raise ValueError(f"unsupported backbone {model_type!r}; known: {sorted(BACKBONES)}")
    return BACKBONES[model_type]


def text_config(path):
    from transformers import AutoConfig

    config = AutoConfig.from_pretrained(path)
    return getattr(config, "text_config", config)


def load_backbone(path, device=None, dtype=torch.bfloat16, attn: str = "sdpa"):
    """Text-only causal LM; refuses a checkpoint with missing text weights (they would be randomly initialised).

    Returns (lm, loading_info)."""
    import transformers

    spec = backbone(path)
    kw = {}
    if device is not None and torch.device(device).type == "cuda":
        kw["device_map"] = {"": str(device)}
    if spec["key_mapping"]:
        kw["key_mapping"] = spec["key_mapping"]
    cls = getattr(transformers, spec["cls"])
    lm, info = cls.from_pretrained(str(path), config=text_config(path), dtype=dtype, attn_implementation=attn,
                                   output_loading_info=True, **kw)
    missing = sorted(k for k in info["missing_keys"] if not k.startswith("lm_head"))
    if missing:
        raise RuntimeError(f"missing weights in {path}: {missing[:5]} ({len(missing)} total)")
    return lm, info


def load_tokenizer(path):
    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained(str(path))
    tok.padding_side = "right"
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
    return tok


def tokenizer_fn(tok):
    """Plain ids: no added special tokens, and special-token strings in post text stay ordinary text."""
    def tokenize(text: str) -> list[int]:
        return tok(text, add_special_tokens=False, split_special_tokens=True)["input_ids"]
    return tokenize


def bos_id(tok) -> int | None:
    """Gemma degrades badly without BOS; ``encode`` prepends exactly one."""
    return tok.bos_token_id


def add_lora(module, r: int, alpha: int, dropout: float):
    from peft import LoraConfig, get_peft_model

    return get_peft_model(module, LoraConfig(r=r, lora_alpha=alpha, lora_dropout=dropout,
                                             target_modules=LORA_TARGETS, bias="none"))


def pad(rows: list[list[int]], pad_id: int, device):
    """Right-padded ids, attention mask and each row's last real index."""
    width = max(len(r) for r in rows)
    ids = torch.full((len(rows), width), pad_id, dtype=torch.long)
    mask = torch.zeros((len(rows), width), dtype=torch.long)
    for i, r in enumerate(rows):
        ids[i, :len(r)] = torch.tensor(r, dtype=torch.long)
        mask[i, :len(r)] = 1
    last = mask.sum(1) - 1
    return ids.to(device), mask.to(device), last.to(device)


class HeadsModel(nn.Module):
    """LoRA text decoder; pooled final-norm hidden state -> one linear head per field."""

    def __init__(self, backbone, hidden: int, sizes: list[int], pool: str = "last", dropout: float = 0.1):
        super().__init__()
        if pool not in POOLS:
            raise ValueError(f"pool must be one of {POOLS}")
        self.backbone = backbone
        self.pool = pool
        self.dropout = nn.Dropout(dropout)
        self.heads = nn.ModuleList([nn.Linear(hidden, k) for k in sizes])

    def pooled(self, ids, mask, last) -> torch.Tensor:
        h = self.backbone(input_ids=ids, attention_mask=mask, use_cache=False).last_hidden_state
        if self.pool == "last":
            return h[torch.arange(h.shape[0], device=h.device), last].float()
        m = mask.unsqueeze(-1).to(h.dtype)
        return ((h * m).sum(1) / m.sum(1).clamp_min(1)).float()

    def forward(self, ids, mask, last) -> list[torch.Tensor]:
        pooled = self.dropout(self.pooled(ids, mask, last))
        return [head(pooled) for head in self.heads]


def step_loss(logits: list[torch.Tensor], labels: torch.Tensor, counts: list[int], weights: list[float] | None = None,
              class_weights: list[torch.Tensor | None] | None = None):
    """Sum over fields of weight x (summed CE over this micro-batch's unmasked rows) / (step's unmasked count).

    ``labels`` is [rows, fields] with -1 for a masked target. Summing this over a step's micro-batches gives, per
    field, the mean over the step's unmasked rows, whatever the split. ``class_weights`` (per field a [K] tensor or
    None) multiplies each row's CE by its gold option's weight. Returns (loss, per-field CE sums)."""
    weights = weights or [1.0] * len(logits)
    class_weights = class_weights or [None] * len(logits)
    total = logits[0].sum() * 0.0
    per_field = []
    for i, lg in enumerate(logits):
        y = labels[:, i]
        keep = y >= 0
        if bool(keep.any()):
            cw = class_weights[i]
            ce = F.cross_entropy(lg[keep].float(), y[keep], reduction="sum",
                                 weight=None if cw is None else cw.to(device=lg.device, dtype=torch.float32))
            total = total + weights[i] * ce / counts[i]
            per_field.append(float(ce.detach()))
        else:
            per_field.append(0.0)
    return total, per_field


def build(cfg: dict, sizes: list[int], device, *, train: bool, adapter_dir=None, heads_path=None) -> HeadsModel:
    """HeadsModel from a base checkpoint dir (``cfg["model"]``); a saved adapter and heads when given."""
    device = torch.device(device)
    lm, _ = load_backbone(cfg["model"], device, dtype=getattr(torch, cfg.get("dtype", "bfloat16")),
                          attn=cfg.get("attn", "sdpa"))
    text = lm.model
    if train and cfg.get("gradient_checkpointing", True):
        text.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    text.config.use_cache = False
    if adapter_dir:
        from peft import PeftModel

        text = PeftModel.from_pretrained(text, str(adapter_dir), is_trainable=train)
    else:
        text = add_lora(text, cfg["lora_r"], cfg["lora_alpha"], cfg["lora_dropout"])
    hidden = lm.config.hidden_size
    del lm  # drops the tied LM head; the text decoder lives on in ``text``
    model = HeadsModel(text, hidden, sizes, cfg.get("pool", "last"), cfg.get("head_dropout", 0.1))
    if heads_path:
        load_heads(model, heads_path)
    model = model.to(device)
    # from_pretrained returns modules in eval mode, and gradient checkpointing only runs in training mode.
    model.train(train)
    return model


# ---------------------------------------------------------------------- checkpoints (classifier-plan §3.3)
def heads_state(model: HeadsModel, fields: list[str]) -> dict[str, torch.Tensor]:
    sd = model.heads.state_dict()
    out = {}
    for i, f in enumerate(fields):
        out[f"{f}.weight"] = sd[f"{i}.weight"].detach().cpu().contiguous()
        out[f"{f}.bias"] = sd[f"{i}.bias"].detach().cpu().contiguous()
    return out


def save_heads(model: HeadsModel, fields: list[str], path: Path) -> None:
    from safetensors.torch import save_file

    save_file(heads_state(model, fields), str(path), metadata={"fields": json.dumps(fields), "pool": model.pool})


def load_heads(model: HeadsModel, path) -> list[str]:
    """Load ``heads.safetensors``; its field list and pooling must match the model."""
    from safetensors import safe_open
    from safetensors.torch import load_file

    with safe_open(str(path), framework="pt") as fh:
        meta = fh.metadata() or {}
    fields = json.loads(meta["fields"])
    if meta.get("pool", model.pool) != model.pool:
        raise ValueError(f"heads were trained with pool={meta.get('pool')}, model has pool={model.pool}")
    if len(fields) != len(model.heads):
        raise ValueError(f"{len(fields)} heads in {path}, model has {len(model.heads)}")
    tensors = load_file(str(path))
    sd = {}
    for i, f in enumerate(fields):
        sd[f"{i}.weight"] = tensors[f"{f}.weight"]
        sd[f"{i}.bias"] = tensors[f"{f}.bias"]
    model.heads.load_state_dict(sd)
    return fields


def rng_state() -> dict:
    out = {"python": random.getstate(), "torch": torch.get_rng_state()}
    if torch.cuda.is_available():
        out["cuda"] = torch.cuda.get_rng_state_all()
    return out


def set_rng_state(state: dict) -> None:
    random.setstate(state["python"])
    torch.set_rng_state(state["torch"])
    if "cuda" in state and torch.cuda.is_available():
        torch.cuda.set_rng_state_all(state["cuda"])


def save_checkpoint(model: HeadsModel, target: Path, fields: list[str], state: dict | None = None,
                    trainer: dict | None = None) -> None:
    """Atomic directory write: adapter/, heads.safetensors, state.json and (resume points) trainer.pt."""
    target = Path(target)
    tmp = target.with_name(target.name + ".tmp")
    shutil.rmtree(tmp, ignore_errors=True)
    tmp.mkdir(parents=True)
    model.backbone.save_pretrained(str(tmp / "adapter"))
    save_heads(model, fields, tmp / "heads.safetensors")
    if trainer is not None:
        torch.save(trainer, tmp / "trainer.pt")
    (tmp / "state.json").write_text(json.dumps(state or {}, indent=1))
    old = target.with_name(target.name + ".old")
    shutil.rmtree(old, ignore_errors=True)
    if target.exists():
        os.replace(target, old)
    os.replace(tmp, target)
    shutil.rmtree(old, ignore_errors=True)


@torch.no_grad()
def predict(model: HeadsModel, encoded: list[list[int]], pad_id: int, batches: list[list[int]], device,
            field_names: list[str]) -> dict[int, dict[str, list[float]]]:
    """Raw logits per field for each encoded row index (dropout off)."""
    device = torch.device(device)
    was_training = model.training
    model.eval()
    out = {}
    amp = dict(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda")
    for b in batches:
        ids, mask, last = pad([encoded[i] for i in b], pad_id, device)
        with torch.autocast(**amp):
            logits = model(ids, mask, last)
        values = [lg.float().cpu().tolist() for lg in logits]
        for j, i in enumerate(b):
            out[i] = {f: values[k][j] for k, f in enumerate(field_names)}
    if was_training:
        model.train()
    return out