"""Stage B: System-1 decision-head SFT (complete-question cross entropy).

Recipe mirrors the NanoJev `hard_lr1e5` run:
  * hard (one-hot) targets on the Choice head;
  * backbone lr 1e-5, decision-head lr 1e-4;
  * mixed question families, complete-question CE (softmax over the whole
    candidate set, not per-leaf binary);
  * optional LoRA on the backbone.

Vision: every sample carries its frame `grid`; the backbone's vision tower
encodes it on EVERY leaf - the same path used at inference. A stage-A
checkpoint (stage_a.pt) is loaded when present so the student's perception
starts aligned with the 9B teacher.

Boolean and Score heads train with CE on trainable logits (the old detached-
tensor BCE was a silent no-op for gradients - fixed here).

Absorbed from Laya: candidate ORDER is permuted per row per epoch (their
build_sequence has the same `option_order` hook). The set-attention Choice
head is order-invariant in theory, but the training distribution is not -
the label candidate always sits where build_dataset put it - so permutation
removes a silent position shortcut. Stage C (calibrate_rlcd.py) then makes
the probabilities honest.
"""

from __future__ import annotations

import argparse
import json
import math


def _oom_skip(exc: Exception, tag: str) -> None:
    """Row-level OOM guard: release the partial graph and keep training.
    Real-game rows vary wildly (64x64 grids, long propositions); one
    pathological row must not kill a multi-hour run."""
    import torch
    torch.cuda.empty_cache()
    print(f"[oom-skip] {tag}: {str(exc).splitlines()[0]}", flush=True)


def _row_sig(row: dict) -> str:
    g = row.get("grid") or []
    return (f"state={len(row.get('state') or '')}ch grid={len(g)}x{len(g[0]) if g and len(g) else 0} "
            f"cands={len(row.get('candidates') or []) or '-'} "
            f"prop={len(row.get('proposition') or '') or '-'} "
            f"desc={len(row.get('description') or '') or '-'}")


def _grid_too_big(row: dict, cap: int = 64) -> bool:
    """Guard kept for tiny-GPU fallbacks; 64x64 trains fine on T4 with
    ckpt + capped candidates (measured 4.4 GiB steady). The step-25 blowup
    was inf/nan loss rows corrupting SDPA numerics, not grid size."""
    g = row.get("grid") or []
    return bool(g) and (len(g) > cap or len(g[0]) > cap)


def _near_oom(row: dict, tag: str) -> bool:
    """Hard gate (v0.4.7): the old T4 run starved to death churning through
    OOM attempts at 14.9GiB reserved — each doomed row burned a full forward
    before failing. Now: above the gate, skip BEFORE the backward (ms-level
    skip instead of a 2-5s doomed forward) and release cache immediately.
    Returns True when the caller must skip this row."""
    import torch
    if torch.cuda.memory_allocated() > 12.5 * 2**30:
        torch.cuda.empty_cache()
        print(f"[gate-skip] gpu={torch.cuda.memory_allocated()/2**30:.2f}GiB "
              f"{tag} | {_row_sig(row)}", flush=True)
        return True
    return False
import random
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import yaml  # noqa: E402


def iter_batches(path: Path, batch_size: int):
    batch = {"choice": [], "boolean": [], "score": []}
    with open(path) as fh:
        for line in fh:
            row = json.loads(line)
            batch[row["family"]].append(row)
            if len(batch["choice"]) >= batch_size:
                yield batch
                batch = {"choice": [], "boolean": [], "score": []}
    yield batch


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/train.yaml")
    args = ap.parse_args()
    cfg = yaml.safe_load(open(args.config))
    b = cfg["stage_b_sft"]
    s1_cfg = cfg.get("system1", {})

    import torch
    import torch.nn.functional as F

    from qwenjev.system1 import QwenJevModel, System1Config
    from qwenjev.vision import TinyVLMConfig

    model_cfg = System1Config(
        backend=s1_cfg.get("backend", "tiny"),
        hf_model=s1_cfg.get("hf_model", "Qwen/Qwen3-VL-2B-Instruct"),
        max_leaf_tokens=int(s1_cfg.get("max_leaf_tokens", 512)),
        score_levels=int(s1_cfg.get("score_levels", 10)),
        leaf_batch_size=int(s1_cfg.get("leaf_batch_size", 32)),
        hf_quant4=bool(s1_cfg.get("hf_quant4", False)),
        pool_vision=bool(s1_cfg.get("pool_vision", False)),
        tiny=TinyVLMConfig(**(s1_cfg.get("tiny") or {})),
    )
    model = QwenJevModel(model_cfg).build()
    torch = model.torch

    # v0.4.13: continue from existing decision heads when the out_dir carries
    # a previous qwenjev_s1.pt (online-teach fine-tunes must keep the teacher
    # generation's calibrated heads; fresh heads would throw away that state)
    prev_heads = Path(cfg["out_dir"]) / "qwenjev_s1.pt"
    if prev_heads.exists():
        model.load_heads(str(cfg["out_dir"]))
        print(f"heads continued -> {prev_heads}")

    # ---- stage A perception alignment (optional but recommended) --------
    stage_a_path = Path(cfg["out_dir"]) / "stage_a.pt"
    if stage_a_path.exists() and model_cfg.backend == "tiny":
        payload = torch.load(stage_a_path, map_location=model.device,
                             weights_only=False)
        if payload.get("student"):
            model._backbone.load_state_dict(payload["student"])
            print(f"stage A loaded -> {stage_a_path}")

    # gradient checkpointing: the flag lives on the base modules, so it
    # stays effective after peft wrapping; REQUIRED for real-game grids on T4
    if s1_cfg.get("grad_ckpt", False) and model_cfg.backend == "hf":
        from qwenjev.hf_compat import enable_grad_ckpt
        if enable_grad_ckpt(model._backbone):
            print("gradient checkpointing ON")

    backbone = model._backbone
    head_params = (
        list(model._choice_proj.parameters())
        + list(model._set_attn.parameters())
        + list(model._boolean_head.parameters())
        + list(model._score_head.parameters())
    )
    peft_model = None
    if b.get("lora", {}).get("enabled") and model_cfg.backend == "hf":
        # unified PEFT path: kbit-prepare for QLoRA bases, fp32-upcast adapters,
        # pretrained vision tower frozen (exclude_vision)
        from qwenjev.hf_compat import maybe_lora, resume_adapters

        # stage A continuity: resume the distillation adapters when present
        # so stage B CONTINUES the same adapters (cold-start LoRA otherwise)
        sa_lora = Path(cfg["out_dir"]) / "stage_a_lora"
        if sa_lora.exists():
            peft_model, trainable = resume_adapters(backbone, str(sa_lora))
            print(f"stage A LoRA resumed -> {sa_lora}")
        else:
            peft_model, trainable = maybe_lora(backbone, b["lora"])
        # rebind onto QwenJevModel (same requirement as the cold-start path:
        # forward must run the WRAPPED model, not the un-frozen base)
        model._backbone = peft_model
    elif b.get("lora", {}).get("enabled"):
        from peft import LoraConfig, get_peft_model

        lcfg = LoraConfig(
            r=int(b["lora"]["r"]), lora_alpha=int(b["lora"]["alpha"]),
            target_modules=list(b["lora"]["target"]), lora_dropout=0.05,
        )
        # CRITICAL: rebind the peft-wrapped model onto QwenJevModel. Without
        # this, forward passes (model._pooled_hf -> self._backbone) run the
        # UNWRAPPED base whose params all require grad; their .grad buffers
        # accumulate outside the optimizer and OOM the T4 in ~25 steps.
        # (masked before v2: the 4-bit kbit path froze the base for us)
        model._backbone = backbone = get_peft_model(backbone, lcfg)
        trainable = [p for p in backbone.parameters() if p.requires_grad]
    else:
        trainable = list(backbone.parameters())

    optim = torch.optim.AdamW(
        [{"params": trainable, "lr": float(b["backbone_lr"])},
         {"params": head_params, "lr": float(b["head_lr"])}],
        weight_decay=0.01,
    )
    # from_pretrained leaves the model in EVAL mode; gradient checkpointing
    # only fires when self.training is True. Without this line the 64x64
    # real-game backward holds ~10 GiB of full activations and OOMs the T4
    # mid-run (diag2/3/4 always called train_mode(True) and never blew up).
    model.train_mode(True)
    print(f"[stage B] trainable: {sum(p.numel() for p in trainable)/1e6:.1f}M "
          f"(backbone {sum(p.numel() for p in backbone.parameters() if p.requires_grad)/1e6:.1f}M + "
          f"heads {sum(p.numel() for p in head_params)/1e6:.1f}M)", flush=True)
    data = Path(cfg["data"]["dataset_dir"]) / "train" / "samples.jsonl"
    if not data.exists():
        raise SystemExit(f"dataset missing: {data}; run build-data first")

    steps = 0
    skipped_big = 0
    epochs = int(b["epochs"])
    grad_acc = max(1, int(b.get("grad_accum", 4)))
    # v0.4.6 crash-resistance: save WIP every N steps. The old T4 died at
    # step ~58/180 and EVERYTHING was lost because saving happened only at
    # the end. With this, worst-case loss = save_every_steps * step_time.
    save_every = int(b.get("save_every_steps", 0) or 0)
    out_dir = Path(cfg["out_dir"])
    max_steps = int(b.get("max_steps", 0) or 0)   # 0 = unlimited
    aux_cap = max(1, int(b.get("aux_rows_per_batch", 8)))
    rng = random.Random(int(cfg.get("data", {}).get("seed", 13)))
    for epoch in range(epochs):
        if max_steps and steps >= max_steps:
            break
        for batch in iter_batches(data, int(b["batch_questions"])):
            losses = []
            # NOTE: every row's loss is backwarded IMMEDIATELY. The hf
            # fp16/QLoRA forward keeps each row's activation graph alive;
            # accumulating whole-batch graphs OOMs even a chunked T4 run.
            # Gradients accumulate across rows/chunks; step once per batch.
            # ---- Choice: complete-question cross entropy ----------------
            for row in batch["choice"]:
                from qwenjev.candidates import CandidateAction, cap_choice_row

                if _grid_too_big(row):
                    skipped_big += 1
                    continue
                # negative-sample wide rows (31 cands x chunked graphs OOMed
                # a 16 GB T4); label index is remapped to the capped list
                try:
                    cand_texts, label = cap_choice_row(
                        row["candidates"], int(row["label"]),
                        int(b.get("max_candidates", 0) or 0), rng)
                    cands = [CandidateAction("CAND", text=t) for t in cand_texts]
                    # order permutation (Laya option_order): track the label
                    order = list(range(len(cands)))
                    rng.shuffle(order)
                    cands = [cands[i] for i in order]
                    label = order.index(label)
                    leaves = model.encode_question(
                        row["state"], cands, row["instruction"], row.get("grid"))
                    if model_cfg.backend == "tiny":
                        out, seq_lens, has_vis = model._pooled_tiny(leaves)
                        torch_ = model.torch
                        vis_n = model._backbone.n_vision_tokens if has_vis else 0
                        pooled = model._backbone.pooled(
                            out, seq_lens + vis_n * torch_.ones_like(seq_lens))
                    else:
                        pooled = model._pooled_hf(leaves)
                    proj = model._choice_proj(pooled)
                    attn_out = model._set_attn(proj.unsqueeze(0)).squeeze(0)
                    logits = attn_out.sum(-1)
                    ce = torch.nn.functional.cross_entropy(
                        logits.unsqueeze(0),
                        torch.tensor([label], device=model.device),
                    )
                    if not torch.isfinite(ce):
                        print(f"[nan-skip] choice row loss={ce.item()} "
                              f"step~{steps+1}", flush=True)
                        continue
                    losses.append(float(ce.detach()))
                    if _near_oom(row, "pre-bwd choice"):
                        optim.zero_grad(set_to_none=True)
                        continue
                    (ce / grad_acc).backward()
                except torch.OutOfMemoryError as exc:
                    _oom_skip(exc, f"choice row (step~{steps+1}) {_row_sig(row)}")
                    optim.zero_grad(set_to_none=True)
                    continue
            # ---- Boolean / Score: trainable-head calibration -------------
            for row in batch["boolean"][:aux_cap]:
                if _grid_too_big(row):
                    skipped_big += 1
                    continue
                try:
                    logits = model.forward_boolean_logits(
                        row["state"], row["instruction"], row["proposition"],
                        row.get("grid"))
                    ce = F.cross_entropy(
                        logits, torch.tensor([row["label"]], device=model.device))
                    if not torch.isfinite(ce):
                        print(f"[nan-skip] boolean row loss={ce.item()} "
                              f"step~{steps+1}", flush=True)
                        continue
                    losses.append(float(ce.detach()))
                    if _near_oom(row, "pre-bwd boolean"):
                        optim.zero_grad(set_to_none=True)
                        continue
                    (ce / grad_acc).backward()
                except torch.OutOfMemoryError as exc:
                    _oom_skip(exc, f"boolean row (step~{steps+1}) {_row_sig(row)}")
                    optim.zero_grad(set_to_none=True)
                    continue
            for row in batch["score"][:aux_cap]:
                if _grid_too_big(row):
                    skipped_big += 1
                    continue
                try:
                    logits = model.forward_score_logits(
                        row["state"], row["instruction"], row["description"],
                        row.get("grid"))
                    ce = F.cross_entropy(
                        logits, torch.tensor([row["label"]], device=model.device))
                    if not torch.isfinite(ce):
                        print(f"[nan-skip] score row loss={ce.item()} "
                              f"step~{steps+1}", flush=True)
                        continue
                    losses.append(float(ce.detach()))
                    if _near_oom(row, "pre-bwd score"):
                        optim.zero_grad(set_to_none=True)
                        continue
                    (ce / grad_acc).backward()
                except torch.OutOfMemoryError as exc:
                    _oom_skip(exc, f"score row (step~{steps+1}) {_row_sig(row)}")
                    optim.zero_grad(set_to_none=True)
                    continue

            if not losses:                       # batch without choice rows
                continue
            torch.nn.utils.clip_grad_norm_(
                [p for group in optim.param_groups for p in group["params"]], 1.0)
            optim.step()
            optim.zero_grad()
            steps += 1
            if steps % 10 == 0:
                torch.cuda.empty_cache()   # v0.4.7: fight allocator fragmentation
            if steps % int(cfg.get("log_every", 10)) == 0:
                print(f"epoch={epoch} step={steps} "
                      f"loss={sum(losses)/len(losses):.4f} "
                      f"gpu={torch.cuda.memory_allocated()/2**30:.2f}GiB "
                      f"big-skip={skipped_big}", flush=True)
            if save_every and steps % save_every == 0:
                # WIP checkpoint: heads via model.save, LoRA via peft.
                # NEVER fatal — a failed save must not kill training.
                try:
                    model.save(str(out_dir))
                    if peft_model is not None:
                        peft_model.save_pretrained(str(out_dir / "hf_lora"))
                    (out_dir / "stage_b_meta.json").write_text(json.dumps(
                        {"steps": steps, "epochs": epoch, "wip": True},
                        indent=2))
                    print(f"[stage B] WIP save @ step {steps}", flush=True)
                except Exception as exc:
                    print(f"[stage B] WIP save failed (non-fatal): {exc}",
                          flush=True)
            if max_steps and steps >= max_steps:
                print(f"max_steps={max_steps} reached; stopping early")
                break

    model.save(cfg["out_dir"])
    # hf backend: persist LoRA adapters so the SFT result is reproducible
    if peft_model is not None:
        try:
            peft_model.save_pretrained(str(Path(cfg["out_dir"]) / "hf_lora"))
        except Exception as exc:
            print(f"[stage B] lora save skipped: {exc}")
    (Path(cfg["out_dir"]) / "stage_b_meta.json").write_text(json.dumps(
        {"steps": steps, "epochs": epochs}, indent=2))
    print(f"stage B done -> {cfg['out_dir']}")


if __name__ == "__main__":
    main()