"""Qwen3.5 text backbone with LoRA and one linear head per field (needs torch; runs on the pod).

Adapted from the read-only reference ~/yt_live_analysis/training/model.py (arm H). Only the text decoder is
loaded: the checkpoints are Qwen3_5ForConditionalGeneration, and loading Qwen3_5ForCausalLM with the
text config maps the language-model weights and ignores the vision tower and the MTP layer. The LM head is
dropped. Rows are right-padded (left padding would run pad tokens through the Gated DeltaNet state before
the post) and each row is pooled at its last real token. The loss is normalised per optimiser step
(``data.step_counts``), so splitting a step into micro-batches leaves its gradient unchanged.
"""
from __future__ import annotations

from pathlib import Path

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

LORA_TARGETS = ['q_proj', 'k_proj', 'v_proj', 'o_proj', 'in_proj_qkv', 'in_proj_z', 'out_proj',
                'gate_proj', 'up_proj', 'down_proj']


def text_config(path):
    from transformers import AutoConfig
    config = AutoConfig.from_pretrained(path)
    return getattr(config, 'text_config', config)


def load_backbone(path, device, dtype=torch.bfloat16, attn='sdpa'):
    """Text-only load; refuses a checkpoint with missing weights (they would be randomly initialised)."""
    from transformers import Qwen3_5ForCausalLM
    kw = {'device_map': {'': str(device)}} if device.type == 'cuda' else {}
    lm, info = Qwen3_5ForCausalLM.from_pretrained(path, config=text_config(path), dtype=dtype, attn_implementation=attn,
                                                  output_loading_info=True, **kw)
    if info['missing_keys']:
        raise RuntimeError('Missing weights in %s: %s' % (path, sorted(info['missing_keys'])[:5]))
    return lm


def load_tokenizer(path):
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(path, split_special_tokens=True)
    tok.padding_side = 'right'
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
    return tok


def tokenizer_fn(tok):
    """Plain token ids: no added special tokens, and special-token strings in the 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 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 backbone; the final-norm hidden state at the last real token feeds one linear head per field."""

    def __init__(self, backbone, hidden: int, sizes: list[int], dropout: float = 0.1):
        super().__init__()
        self.backbone = backbone
        self.dropout = nn.Dropout(dropout)
        self.heads = nn.ModuleList([nn.Linear(hidden, k) for k in sizes])

    def forward(self, ids, mask, last) -> list[torch.Tensor]:
        h = self.backbone(input_ids=ids, attention_mask=mask, use_cache=False).last_hidden_state
        pooled = h[torch.arange(h.shape[0], device=h.device), last].float()
        pooled = self.dropout(pooled)
        return [head(pooled) for head in self.heads]


def step_loss(logits: list[torch.Tensor], labels: torch.Tensor, counts: list[int]):
    """Sum over fields of (summed cross-entropy over this micro-batch's non-null rows) / (step's non-null count).

    ``labels`` is [rows, fields] with -1 for a null target. Summing this over a step's micro-batches gives,
    per field, the mean over the step's non-null rows, whatever the split. Returns (loss, per-field sums).
    """
    total = logits[0].sum() * 0.0
    per_field = []
    for i, lg in enumerate(logits):
        y = labels[:, i]
        keep = y >= 0
        if bool(keep.any()):
            ce = F.cross_entropy(lg[keep].float(), y[keep], reduction='sum')
            total = total + 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):
    lm = load_backbone(cfg['model'], device, dtype=getattr(torch, cfg.get('dtype', 'bfloat16')))
    text = lm.model
    if train:
        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'])
    lm.lm_head = None
    model = HeadsModel(text, lm.config.hidden_size, sizes, cfg['head_dropout'])
    if heads_path:
        model.heads.load_state_dict(torch.load(heads_path, map_location='cpu'))
    model = model.to(device)
    # from_pretrained returns modules in eval mode, and gradient checkpointing only runs in training mode.
    model.train(train)
    return model


def save_checkpoint(model: HeadsModel, target: Path, extra: dict | None = None, trainer=None) -> None:
    """Atomic directory write: adapter/, heads.pt, state.json and (for resume points) trainer.pt."""
    import json
    import shutil
    tmp = target.with_name(target.name + '.tmp')
    shutil.rmtree(tmp, ignore_errors=True)
    tmp.mkdir(parents=True)
    model.backbone.save_pretrained(str(tmp / 'adapter'))
    torch.save(model.heads.state_dict(), tmp / 'heads.pt')
    if trainer is not None:
        torch.save(trainer, tmp / 'trainer.pt')
    (tmp / 'state.json').write_text(json.dumps(extra or {}))
    shutil.rmtree(target, ignore_errors=True)
    tmp.rename(target)


@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)."""
    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