#!/usr/bin/env python3
"""
sfd_4b.py (formerly 16d_train_filtered_distill.py)

Rejection-sampled cross-size distillation: same teacher/student setup as
16c, but TRAINING IS FILTERED. We sample y_teacher from the teacher (with
privileged demo), parse its predicted bbox, and only update the student
when the teacher itself was correct (predicted center inside GT bbox).

Per micro-batch:
  1. Teacher.generate(teacher_prompt_with_demo) -> y_teacher (no grad)
  2. Parse y_teacher's response -> predicted bbox
  3. If pred bbox center NOT in GT bbox: skip (zero loss, no grad)
  4. Else: student fwd on [student_prompt, y_teacher] -> s_logits
           teacher fwd on [teacher_prompt, y_teacher] -> t_logits
           reverse KL on y_teacher positions -> backward

Why: filters teacher noise. The 4B teacher only gets ~34% click-accuracy
on this dataset; the other 66% of teacher rollouts are wrong. Distilling
those wrong rollouts wastes student capacity. Rejection sampling means
fewer effective updates per micro-batch but cleaner per-update signal.

Differs from 16c (on-policy from student) by being OFF-POLICY (uses
teacher rollouts) AND quality-filtered.

Originally derived from 16c, which differs from 16b:
  - Teacher loaded from --teacher-model (e.g. Qwen3.5-4B), NOT a copy of
    student. Vocab/tokenizer must match (true across all Qwen3.5-* sizes).
  - Teacher is FIXED (no EMA updates). The student's distribution is pulled
    toward a stronger, demo-conditioned teacher.
  - Demo suffix is "Now answer with the JSON only:" (NOT the paper's CoT-
    encouraging "...including the thinking process") — keeps the teacher
    distribution terse, which is what we want to distill into the student.
  - Both prompts force enable_thinking=False on the chat template.
  - **No QLoRA / no PEFT**: full-parameter bf16 student.
  - EMA teacher state is a full bf16 state-dict shadow on GPU (~4 GB for 2 B
    params). Reuses the ParamSwap pattern unchanged — _iter_trainable already
    filters by requires_grad, so it auto-covers full-FT.
  - EMA is REQUIRED (--ema-alpha must be > 0). Without LoRA there is no
    disable_adapter() fallback to obtain a teacher distinct from the student.
  - Gradient checkpointing on; use_cache=False on the model config.

Training loop per step (unchanged from 16):
  1. Student.generate(prompt_only) -> y_student
  2. Student forward WITH GRAD on [prompt, y_student] -> s_logits
  3. Teacher forward NO GRAD via ParamSwap(EMA shadow) on
        [prompt_with_demo, y_student] -> t_logits
  4. KL(student || teacher) on y_student positions.
  5. EMA update phi <- alpha * theta + (1 - alpha) * phi per paper Line 19.

Memory footprint on H100 80 GB (Qwen3.5-2B, vision frozen, bf16):
    params bf16             ~3.0 GB   (LM only, vision frozen)
    grads bf16              ~3.0 GB
    AdamW states bf16       ~6.0 GB
    EMA shadow bf16         ~3.0 GB
    activations (grad-ckpt) ~8-12 GB  (bs=2, seq=2048)
    KV cache during generate ~1-2 GB
    ---------------------------------
    total                   ~25-30 GB  (comfortable headroom)

Usage:
    python -m windom train-sfd4b \\
        --model Qwen/Qwen3.5-2B \\
        --data-dir data/win11react_v42_500/gemma4_ready \\
        --output-dir checkpoints/v42_4b_hint \\
        --epochs 2 --batch_size 2 --grad_accum 16 --lr 5e-6 \\
        --kl-type reverse --kl-temperature 1.0 \\
        --on-policy-max-new-tokens 40 \\
        --save-steps 200 --wandb
"""
import argparse
import json
import os
import threading
from pathlib import Path

import torch


# ─── Background HF upload helper (non-blocking) ────────────────────────────
_upload_threads: "list[threading.Thread]" = []

def _bg_upload(ck_path, repo_id, path_in_repo, commit_msg, token):
    try:
        from huggingface_hub import HfApi
        api = HfApi(token=token)
        api.create_repo(repo_id, exist_ok=True, repo_type="model")
        api.upload_folder(folder_path=str(ck_path), repo_id=repo_id,
                          path_in_repo=path_in_repo, commit_message=commit_msg)
        print(f"  → BG upload done: {repo_id}/{path_in_repo}", flush=True)
    except Exception as e:
        print(f"  WARN BG upload failed for {path_in_repo}: {e}", flush=True)

def schedule_bg_upload(ck_path, repo_id, path_in_repo, commit_msg):
    t = threading.Thread(target=_bg_upload,
                         args=(ck_path, repo_id, path_in_repo, commit_msg,
                               os.environ.get("HF_TOKEN")), daemon=True)
    t.start()
    _upload_threads.append(t)

def join_bg_uploads(timeout_per: int = 600):
    if not _upload_threads: return
    print(f"Waiting for {len(_upload_threads)} BG upload(s) (timeout {timeout_per}s each)...",
          flush=True)
    for t in _upload_threads:
        t.join(timeout=timeout_per)
        if t.is_alive():
            print(f"  WARN upload thread still running after timeout", flush=True)
import torch.nn as nn
import torch.nn.functional as F

# Shared helpers (dataset loader, prompts, parser, scoring, wandb).
from windom.core import (
    GROUNDING_SYSTEM_PROMPT_QWEN,
    GROUNDING_TYPES,
    _to_structured_content,
    build_dataset,
    init_wandb,
    parse_tool_call,
    point_in_bbox,
)


# ─── EMA helpers (two-model variant: separate student + teacher) ────────────

@torch.no_grad()
def copy_student_into_teacher(student, teacher):
    """Hard-copy ALL student params into teacher (used at startup)."""
    s_state = student.state_dict()
    t_state = teacher.state_dict()
    for n, p_s in s_state.items():
        if n in t_state:
            t_state[n].data.copy_(p_s.data)

@torch.no_grad()
def ema_update_teacher(student, teacher, alpha):
    """phi <- (1 - alpha) * phi + alpha * theta, only over student-trainable
    params (vision tower stays frozen and identical in both models)."""
    s_named = dict(student.named_parameters())
    for n, p_t in teacher.named_parameters():
        p_s = s_named.get(n)
        if p_s is None or not p_s.requires_grad:
            continue
        p_t.data.mul_(1 - alpha).add_(p_s.data, alpha=alpha)


# ─── Teacher prompt construction ────────────────────────────────────────────

DEMO_PREFIX = "\nThis is an example for a response to the question:\n"
DEMO_SUFFIX = "\nNow answer with the tool call only:"
STUDENT_NO_REASON_SUFFIX = (
    "\nOutput ONLY the tool call. Do NOT reason, explain, or add anything else."
)

# --teacher-mode hint: jitter applied to GT bbox coords (0-1000 scale)
HINT_JITTER = 30


def build_teacher_messages(student_messages, teacher_mode: str = "example",
                            gt_bbox_xyxy=None,
                            official_prompt: bool = False,
                            image_w: int = 1000, image_h: int = 1000):
    """Insert GT guidance into the teacher's user turn.

    teacher_mode='example': full GT JSON shown verbatim (original behaviour).
    teacher_mode='hint':    approximate bbox shown as a hint with ±HINT_JITTER
                            noise, forcing the teacher to do real visual grounding.

    Under official_prompt=True, hint-mode bbox coords are scaled to PIXEL
    space (using image_w, image_h) so they're consistent with the
    system-prompt's "screen's resolution is W x H" framing. Example mode
    re-uses the assistant turn's existing target string (already in
    pixel coords when --official-prompt is set on the dataset).
    """
    import random

    # Extract assistant response text (used for 'example' mode and to get raw bbox)
    demo = None
    for m in student_messages:
        if m.get("role") == "assistant":
            for p in m.get("content", []):
                if isinstance(p, dict) and p.get("type") == "text":
                    demo = p.get("text", ""); break
            break
    if not demo:
        return student_messages

    if teacher_mode == "example":
        suffix_text = f"{DEMO_PREFIX}{demo}{DEMO_SUFFIX}"
    else:  # hint: show approx bbox region with jitter
        if gt_bbox_xyxy is not None:
            x1, y1, x2, y2 = gt_bbox_xyxy   # in 0-1000 normalized space
            # Hint coords stay in 0-1000 normalized — Qwen3-VL emits in this
            # range regardless of the system prompt's pixel-resolution claim.
            def jitter(v):
                return int(max(0, min(1000, v + random.randint(-HINT_JITTER, HINT_JITTER))))
            hx1, hy1, hx2, hy2 = jitter(x1), jitter(y1), jitter(x2), jitter(y2)
            suffix_text = (
                f"\nHint: the target element is approximately in the region "
                f"[{hx1}, {hy1}, {hx2}, {hy2}] (xmin, ymin, xmax, ymax).\n"
                f"Now answer with the tool call only:"
            )
        else:
            suffix_text = f"{DEMO_PREFIX}{demo}{DEMO_SUFFIX}"

    out = []
    for m in student_messages:
        if m.get("role") == "user":
            parts = [dict(p) for p in m.get("content", [])]
            parts.append({"type": "text", "text": suffix_text})
            out.append({"role": "user", "content": parts})
        else:
            out.append({"role": m["role"],
                        "content": [dict(p) if isinstance(p, dict) else p
                                    for p in m.get("content", [])]})
    return out


# ─── Collator — builds student and teacher prompt-only batches ──────────────

class OnPolicyCollator:
    """For each feature, produce THREE tokenized variants:
      1. student_prompt — [system + user] only, ends at "<|im_start|>assistant\\n"
      2. student_full   — student_prompt + GT response (for fallback NTP if needed)
      3. teacher_prompt — [system + user_with_demo] only, ends at same position
    Returns a dict with student_* and teacher_* keys plus pixel_values,
    image_grid_thw (shared — same image).
    """
    def __init__(self, processor, max_length: int = 2048, teacher_mode: str = "example",
                  official_prompt: bool = False):
        self.processor = processor
        self.max_length = max_length
        self.teacher_mode = teacher_mode
        self.official_prompt = official_prompt

    def __call__(self, features):
        out = {}
        student_prompt_batches = []
        teacher_prompt_batches = []
        gt_response_batches = []

        for f in features:
            messages = f["messages"]
            img_raw = f["images"][0]
            # HF Datasets' Image() feature yields PIL Image objects directly.
            # Fall back to open() only if we were given a path string.
            from PIL import Image as PILImage
            if hasattr(img_raw, "convert"):
                image = img_raw.convert("RGB")
            else:
                image = PILImage.open(img_raw).convert("RGB")

            # Student messages: prompt only (strip assistant turn)
            student_prompt_msgs = [m for m in messages
                                   if m.get("role") != "assistant"]
            # Build teacher prompt from the ORIGINAL student msgs (no suffix yet)
            # so demo doesn't get sandwiched with student-only instruction.
            import copy
            gt_bbox = f.get("bbox", None)
            gt_bbox_xyxy = [gt_bbox[1], gt_bbox[0], gt_bbox[3], gt_bbox[2]] if gt_bbox else None
            img_w = int(f.get("image_w") or image.width)
            img_h = int(f.get("image_h") or image.height)
            teacher_prompt_msgs = build_teacher_messages(
                copy.deepcopy(student_prompt_msgs) + [{
                    "role": "assistant",
                    "content": [p for m in messages
                                if m.get("role") == "assistant"
                                for p in m.get("content", [])],
                }],
                teacher_mode=self.teacher_mode,
                gt_bbox_xyxy=gt_bbox_xyxy,
                official_prompt=self.official_prompt,
                image_w=img_w, image_h=img_h)
            # But we need prompt-only (no assistant) for teacher too
            teacher_prompt_msgs = [m for m in teacher_prompt_msgs
                                   if m.get("role") != "assistant"]
            # NOW append the no-reason suffix only to the STUDENT's user turn.
            # Skip under official prompt — the system prompt already says
            # "respond with ONLY a tool call".
            if not self.official_prompt:
                for m in student_prompt_msgs:
                    if m.get("role") == "user":
                        for c in m.get("content", []):
                            if isinstance(c, dict) and c.get("type") == "text":
                                if STUDENT_NO_REASON_SUFFIX.strip() not in c.get("text", ""):
                                    c["text"] = c.get("text", "") + STUDENT_NO_REASON_SUFFIX
                                break
                        break

            # Get GT response text
            gt_text = ""
            for m in messages:
                if m.get("role") == "assistant":
                    for p in m.get("content", []):
                        if isinstance(p, dict) and p.get("type") == "text":
                            gt_text = p.get("text", ""); break
                    break

            # Force thinking off — Qwen3.5 chat template defaults vary across
            # model sizes; this keeps the rendered prompt consistent.
            s_prompt_text = self.processor.apply_chat_template(
                student_prompt_msgs, tokenize=False, add_generation_prompt=True,
                enable_thinking=False)
            t_prompt_text = self.processor.apply_chat_template(
                teacher_prompt_msgs, tokenize=False, add_generation_prompt=True,
                enable_thinking=False)

            # Student: tokenize with image. truncation disabled — Qwen3-VL
            # uses dynamic image tokenization, so a hard max_length cuts image
            # tokens and breaks the image-count check in the model forward.
            s_enc = self.processor(
                text=[s_prompt_text], images=[image], return_tensors="pt",
                padding=False)
            # Teacher: tokenize with same image
            t_enc = self.processor(
                text=[t_prompt_text], images=[image], return_tensors="pt",
                padding=False)

            # GT response tokens (no specials, no padding)
            tok = getattr(self.processor, "tokenizer", self.processor)
            gt_ids = tok(gt_text, add_special_tokens=False,
                         return_tensors="pt")["input_ids"][0]
            # Append <|im_end|> terminator to GT
            im_end = tok.convert_tokens_to_ids("<|im_end|>")
            if im_end is not None and im_end >= 0 and gt_ids[-1].item() != im_end:
                gt_ids = torch.cat([gt_ids, torch.tensor([im_end])])

            student_prompt_batches.append(s_enc)
            teacher_prompt_batches.append(t_enc)
            gt_response_batches.append(gt_ids)

        # Pad prompts to same length per side (use LEFT padding for generate).
        def _pad_left(enc_list, key, pad_id):
            tensors = [e[key][0] for e in enc_list]
            max_len = max(t.shape[0] for t in tensors)
            padded = []
            for t in tensors:
                pad_len = max_len - t.shape[0]
                if pad_len > 0:
                    pad = torch.full((pad_len,), pad_id, dtype=t.dtype)
                    t = torch.cat([pad, t])
                padded.append(t)
            return torch.stack(padded)

        tok = getattr(self.processor, "tokenizer", self.processor)
        pad_id = tok.pad_token_id if tok.pad_token_id is not None else 0

        out["student_input_ids"] = _pad_left(student_prompt_batches, "input_ids", pad_id)
        out["student_attention_mask"] = _pad_left(
            student_prompt_batches, "attention_mask", 0)
        out["teacher_input_ids"] = _pad_left(teacher_prompt_batches, "input_ids", pad_id)
        out["teacher_attention_mask"] = _pad_left(
            teacher_prompt_batches, "attention_mask", 0)
        # Qwen3.5-VL requires mm_token_type_ids for M-RoPE to disambiguate
        # image vs text token positions. Pad with 0 (text).
        if "mm_token_type_ids" in student_prompt_batches[0]:
            out["student_mm_token_type_ids"] = _pad_left(
                student_prompt_batches, "mm_token_type_ids", 0)
        if "mm_token_type_ids" in teacher_prompt_batches[0]:
            out["teacher_mm_token_type_ids"] = _pad_left(
                teacher_prompt_batches, "mm_token_type_ids", 0)

        # Student and teacher images are IDENTICAL — share pixel_values
        # from the student batch (processor stores them flat).
        def _stack_flat(key):
            all_ = []
            for e in student_prompt_batches:
                if key in e and e[key] is not None:
                    all_.append(e[key])
            return torch.cat(all_, dim=0) if all_ else None
        out["pixel_values"] = _stack_flat("pixel_values")
        out["image_grid_thw"] = torch.stack([
            e["image_grid_thw"][0] for e in student_prompt_batches])

        # Separate teacher pixel_values (same image, re-encoded might differ due to resize)
        def _stack_flat_t(key):
            all_ = []
            for e in teacher_prompt_batches:
                if key in e and e[key] is not None:
                    all_.append(e[key])
            return torch.cat(all_, dim=0) if all_ else None
        out["teacher_pixel_values"] = _stack_flat_t("pixel_values")
        out["teacher_image_grid_thw"] = torch.stack([
            e["image_grid_thw"][0] for e in teacher_prompt_batches])

        # Pad GT responses (RIGHT-pad) for reference.
        gt_max = max(g.shape[0] for g in gt_response_batches)
        gt_padded = []
        gt_mask = []
        for g in gt_response_batches:
            pad_len = gt_max - g.shape[0]
            if pad_len > 0:
                g_ = torch.cat([g, torch.full((pad_len,), pad_id, dtype=g.dtype)])
                m_ = torch.cat([torch.ones_like(g), torch.zeros(pad_len, dtype=torch.long)])
            else:
                g_ = g; m_ = torch.ones_like(g)
            gt_padded.append(g_); gt_mask.append(m_)
        out["gt_response_ids"] = torch.stack(gt_padded)
        out["gt_response_mask"] = torch.stack(gt_mask)

        # 16d: also pass GT bboxes through so train_step can filter on
        # teacher correctness. Dataset stores [y1, x1, y2, x2]; un-swap to
        # [x1, y1, x2, y2] for the click-in-bbox test.
        # Under official_prompt, the teacher emits PIXEL coords; we also
        # build a pixel-space copy of the GT bbox so the rejection check
        # compares predictions and GT in the SAME coordinate space.
        gt_bboxes_xyxy = []
        gt_bboxes_xyxy_pixel = []
        for f in features:
            b = f.get("bbox", [-1.0, -1.0, -1.0, -1.0])
            xyxy_norm = [b[1], b[0], b[3], b[2]]
            gt_bboxes_xyxy.append(xyxy_norm)
            iw = float(f.get("image_w") or 1000.0)
            ih = float(f.get("image_h") or 1000.0)
            gt_bboxes_xyxy_pixel.append([
                xyxy_norm[0] * iw / 1000.0, xyxy_norm[1] * ih / 1000.0,
                xyxy_norm[2] * iw / 1000.0, xyxy_norm[3] * ih / 1000.0,
            ])
        out["gt_bboxes_xyxy"] = gt_bboxes_xyxy           # 0-1000 normalized
        out["gt_bboxes_xyxy_pixel"] = gt_bboxes_xyxy_pixel  # pixel space

        return out


# ─── Training loop ──────────────────────────────────────────────────────────

def reverse_kl(s_logits, t_logits, T=1.0):
    s_logp = F.log_softmax(s_logits / T, dim=-1)
    t_logp = F.log_softmax(t_logits / T, dim=-1)
    s_prob = s_logp.exp()
    return (s_prob * (s_logp - t_logp)).sum(-1).mean() * (T ** 2)

def jsd(s_logits, t_logits, T=1.0, beta=0.5):
    s_logp = F.log_softmax(s_logits / T, dim=-1)
    t_logp = F.log_softmax(t_logits / T, dim=-1)
    s_prob = s_logp.exp(); t_prob = t_logp.exp()
    m_prob = beta * s_prob + (1 - beta) * t_prob
    m_logp = (m_prob + 1e-12).log()
    kl_sm = (s_prob * (s_logp - m_logp)).sum(-1).mean()
    kl_tm = (t_prob * (t_logp - m_logp)).sum(-1).mean()
    return (beta * kl_sm + (1 - beta) * kl_tm) * (T ** 2)


def train_step(model, teacher, batch, *, kl_type, kl_temperature, jsd_beta,
                max_new_tokens, sample_temp, stop_ids, tokenizer):
    """16d rejection-sampled distillation step.

    Returns (loss_or_zero, was_kept). loss_or_zero has requires_grad even
    when zero, so .backward() is safe. was_kept is True if we actually
    contributed a gradient.
    """
    device = next(model.parameters()).device
    dtype = model.dtype

    s_ids   = batch["student_input_ids"].to(device)
    s_mask  = batch["student_attention_mask"].to(device)
    t_ids   = batch["teacher_input_ids"].to(device)
    t_mask  = batch["teacher_attention_mask"].to(device)
    pix     = batch["pixel_values"].to(device, dtype=dtype) if batch.get("pixel_values") is not None else None
    grid    = batch["image_grid_thw"].to(device)
    t_pix   = batch["teacher_pixel_values"].to(device, dtype=dtype) if batch.get("teacher_pixel_values") is not None else None
    t_grid  = batch["teacher_image_grid_thw"].to(device)
    s_mm_ids = batch["student_mm_token_type_ids"].to(device) if "student_mm_token_type_ids" in batch else None
    t_mm_ids = batch["teacher_mm_token_type_ids"].to(device) if "teacher_mm_token_type_ids" in batch else None
    gt_bboxes_xyxy = batch["gt_bboxes_xyxy"]   # python list (0-1000 normalized)
    gt_bboxes_xyxy_pixel = batch.get("gt_bboxes_xyxy_pixel", None)  # pixel space

    B = s_ids.shape[0]

    # ── 1. TEACHER generates y_teacher (no grad) — its own answer with demo ──
    with torch.no_grad():
        gen_kwargs = dict(
            input_ids=t_ids, attention_mask=t_mask,
            pixel_values=t_pix, image_grid_thw=t_grid,
            max_new_tokens=max_new_tokens,
            do_sample=False,                      # greedy: deterministic teacher
            eos_token_id=stop_ids if stop_ids else None,
            pad_token_id=stop_ids[-1] if stop_ids else None,
        )
        if t_mm_ids is not None:
            gen_kwargs["mm_token_type_ids"] = t_mm_ids
        gen = teacher.generate(**gen_kwargs)
    y_teacher = gen[:, t_ids.shape[1]:]     # (B, gen_len)
    gen_len = y_teacher.shape[1]
    if gen_len == 0:
        return torch.zeros((), device=device, requires_grad=True), False

    # ── 2. Parse y_teacher per sample, check correctness vs GT ──────────
    # Both pred and GT are in 0-1000 normalized space (Qwen3-VL native
    # output range, official evaluator divides model output by 1000 too).
    keep_mask = []
    for i in range(B):
        text = tokenizer.decode(y_teacher[i].tolist(), skip_special_tokens=False)
        pred = parse_tool_call(text)
        ok = point_in_bbox(pred, gt_bboxes_xyxy[i])
        keep_mask.append(ok)
    if not any(keep_mask):
        return torch.zeros((), device=device, requires_grad=True), False

    # Filter batch dim — keep only correct samples for the gradient pass.
    keep_idx = [i for i, k in enumerate(keep_mask) if k]
    keep_t = torch.tensor(keep_idx, device=device)

    s_ids_k    = s_ids.index_select(0, keep_t)
    s_mask_k   = s_mask.index_select(0, keep_t)
    t_ids_k    = t_ids.index_select(0, keep_t)
    t_mask_k   = t_mask.index_select(0, keep_t)
    y_teach_k  = y_teacher.index_select(0, keep_t)
    s_mm_k     = s_mm_ids.index_select(0, keep_t) if s_mm_ids is not None else None
    t_mm_k     = t_mm_ids.index_select(0, keep_t) if t_mm_ids is not None else None

    # Image features per kept sample. pixel_values is flat-cat across all
    # samples (Qwen3.5-VL dynamic resolution). Split by patch counts from
    # image_grid_thw, then keep only the filtered indices.
    def _filter_pix(pixel_values, image_grid_thw, keep_t):
        if pixel_values is None:
            return None, image_grid_thw.index_select(0, keep_t)
        patch_counts = [int(g[0] * g[1] * g[2]) for g in image_grid_thw]
        splits = torch.split(pixel_values, patch_counts, dim=0)
        kept_pix = torch.cat([splits[i] for i in keep_t.tolist()], dim=0)
        kept_grid = image_grid_thw.index_select(0, keep_t)
        return kept_pix, kept_grid

    pix_k,   grid_k   = _filter_pix(pix,   grid,   keep_t)
    t_pix_k, t_grid_k = _filter_pix(t_pix, t_grid, keep_t)

    # ── 3. STUDENT forward WITH GRAD on [student_prompt, y_teacher] ─────
    s_full   = torch.cat([s_ids_k, y_teach_k], dim=-1)
    s_full_m = torch.cat([s_mask_k, torch.ones_like(y_teach_k)], dim=-1)
    s_full_mm = None
    if s_mm_k is not None:
        s_full_mm = torch.cat([s_mm_k, torch.zeros_like(y_teach_k)], dim=-1)
    s_kwargs = dict(input_ids=s_full, attention_mask=s_full_m,
                    pixel_values=pix_k, image_grid_thw=grid_k)
    if s_full_mm is not None:
        s_kwargs["mm_token_type_ids"] = s_full_mm
    s_out = model(**s_kwargs)
    s_logits = s_out.logits
    s_gen_logits = s_logits[:, s_ids_k.shape[1] - 1 : -1, :]

    # ── 4. TEACHER forward NO GRAD on [teacher_prompt, y_teacher] ───────
    t_full   = torch.cat([t_ids_k, y_teach_k], dim=-1)
    t_full_m = torch.cat([t_mask_k, torch.ones_like(y_teach_k)], dim=-1)
    t_full_mm = None
    if t_mm_k is not None:
        t_full_mm = torch.cat([t_mm_k, torch.zeros_like(y_teach_k)], dim=-1)
    t_kwargs = dict(input_ids=t_full, attention_mask=t_full_m,
                    pixel_values=t_pix_k, image_grid_thw=t_grid_k)
    if t_full_mm is not None:
        t_kwargs["mm_token_type_ids"] = t_full_mm
    with torch.no_grad():
        t_out = teacher(**t_kwargs)
        t_logits = t_out.logits.detach()
        t_gen_logits = t_logits[:, t_ids_k.shape[1] - 1 : -1, :]

    # ── 5. KL on y_teacher positions ────────────────────────────────────
    if kl_type == "reverse":
        kl = reverse_kl(s_gen_logits, t_gen_logits, T=kl_temperature)
    elif kl_type == "jsd":
        kl = jsd(s_gen_logits, t_gen_logits, T=kl_temperature, beta=jsd_beta)
    else:
        t_logp = F.log_softmax(t_gen_logits / kl_temperature, dim=-1)
        s_logp = F.log_softmax(s_gen_logits / kl_temperature, dim=-1)
        t_prob = t_logp.exp()
        kl = (t_prob * (t_logp - s_logp)).sum(-1).mean() * (kl_temperature ** 2)

    return kl, True


# ─── Main ───────────────────────────────────────────────────────────────────

def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="Qwen/Qwen3.5-2B",
                    help="Student model HF id or local path.")
    ap.add_argument("--teacher-model", default="Qwen/Qwen3.5-4B",
                    help="Teacher model HF id or local path. Must share "
                         "vocabulary with student (same Qwen3.5 family).")
    ap.add_argument("--subfolder", default="",
                    help="HF subfolder for student (e.g. 'checkpoint-200', 'final').")
    ap.add_argument("--teacher-subfolder", default="",
                    help="HF subfolder for teacher (default: empty / repo root).")
    ap.add_argument("--data-dir", type=Path, required=True)
    ap.add_argument("--output-dir", type=Path, required=True)
    ap.add_argument("--epochs", type=int, default=2)
    ap.add_argument("--max-steps", type=int, default=None,
                    help="Stop when global_step reaches this value (use with resume to finish epoch).")
    ap.add_argument("--batch_size", type=int, default=2)
    ap.add_argument("--grad_accum", type=int, default=16)
    ap.add_argument("--lr", type=float, default=5e-6)
    ap.add_argument("--warmup-ratio", type=float, default=0.03)
    ap.add_argument("--max-grad-norm", type=float, default=1.0)
    ap.add_argument("--max-seq-length", type=int, default=2048)
    ap.add_argument("--unfreeze-vision", action="store_true",
                    help="Train the vision tower too (default: frozen).")
    ap.add_argument("--kl-type", choices=["reverse", "forward", "jsd"], default="reverse")
    ap.add_argument("--kl-temperature", type=float, default=1.0)
    ap.add_argument("--jsd-beta", type=float, default=0.5)
    ap.add_argument("--ema-alpha", type=float, default=0.01)
    ap.add_argument("--on-policy-max-new-tokens", type=int, default=40)
    ap.add_argument("--on-policy-temperature", type=float, default=1.0)
    ap.add_argument("--official-prompt", action="store_true",
                    help="Use the official ScreenSpot-Pro prompt format "
                         "(system role + <tools> JSON schema, pixel-coord targets) "
                         "throughout training and teacher rollouts. Matches the "
                         "public-leaderboard evaluator.")
    ap.add_argument("--teacher-mode", choices=["example", "hint"], default="example",
                    help="'example': show GT answer verbatim (default). "
                         "'hint': show approximate bbox with jitter, forcing real grounding.")
    ap.add_argument("--save-steps", type=int, default=200)
    ap.add_argument("--hf-repo", type=str, default=None,
                    help="HuggingFace repo id to push each checkpoint to (e.g. Chengheng/webos-16d-2b).")
    ap.add_argument("--log-every", type=int, default=10)
    ap.add_argument("--max-samples", type=int, default=None)
    ap.add_argument("--wandb", action="store_true")
    ap.add_argument("--wandb-project", type=str, default="webos-gui-grounding")
    ap.add_argument("--wandb-run-name", type=str, default=None)
    args = ap.parse_args(argv)

    # 16c uses a fixed external teacher — EMA is irrelevant. Ignore --ema-alpha.
    if args.ema_alpha != 0:
        print(f"  NOTE: --ema-alpha={args.ema_alpha} ignored (16c teacher is fixed).")

    if args.wandb:
        os.environ["WANDB_PROJECT"] = args.wandb_project
        if args.wandb_run_name:
            os.environ["WANDB_NAME"] = args.wandb_run_name
        try:
            import wandb
            wandb.init(project=args.wandb_project, name=args.wandb_run_name,
                       config=vars(args))
        except Exception as e:
            print(f"  wandb init failed: {e}; continuing without")

    # ── Load full model in bf16 (no quantization, no PEFT) ────────────
    from transformers import AutoProcessor, AutoModelForImageTextToText

    s_kw = {"subfolder": args.subfolder} if args.subfolder else {}
    print(f"Loading STUDENT {args.model}" + (f" (subfolder={args.subfolder})" if args.subfolder else "") + " in bf16 ...")
    processor = AutoProcessor.from_pretrained(args.model, **s_kw)
    model = AutoModelForImageTextToText.from_pretrained(
        args.model,
        torch_dtype=torch.bfloat16,
        device_map="cuda:0",
        **s_kw,
    )
    # Gradient checkpointing requires use_cache=False.
    model.config.use_cache = False
    # use_cache=False is required for gradient checkpointing to fire.
    model.config.use_cache = False
    # Gradient checkpointing on the student trades compute for activation
    # memory. It does NOT change which parameters receive gradients — all
    # trainable params still get full gradients (paper-faithful "full FT").
    model.gradient_checkpointing_enable(
        gradient_checkpointing_kwargs={"use_reentrant": False}
    )
    if hasattr(model, "enable_input_require_grads"):
        model.enable_input_require_grads()

    # Freeze vision tower by default (standard for VL SDFT).
    if not args.unfreeze_vision:
        n_frozen = 0
        for n, p in model.named_parameters():
            if any(tag in n for tag in ("vision", "visual", "patch_embed")):
                p.requires_grad = False
                n_frozen += p.numel()
        print(f"  vision/visual params frozen: {n_frozen/1e6:.1f} M")

    n_train = sum(p.numel() for p in model.parameters() if p.requires_grad)
    n_total = sum(p.numel() for p in model.parameters())
    print(f"  trainable: {n_train/1e9:.3f} B / {n_total/1e9:.3f} B "
          f"({100*n_train/n_total:.1f}%)")

    # ── Load teacher: SEPARATE model from --teacher-model (cross-size) ──
    t_kw = {"subfolder": args.teacher_subfolder} if args.teacher_subfolder else {}
    print(f"Loading FIXED TEACHER from {args.teacher_model}" + (f" (subfolder={args.teacher_subfolder})" if args.teacher_subfolder else "") + " in bf16 ...")
    teacher = AutoModelForImageTextToText.from_pretrained(
        args.teacher_model,
        torch_dtype=torch.bfloat16,
        device_map="cuda:0",
        **t_kw,
    )
    teacher.config.use_cache = False
    for p in teacher.parameters():
        p.requires_grad = False
    teacher.eval()
    n_t = sum(p.numel() for p in teacher.parameters())
    print(f"  teacher loaded, {n_t/1e9:.2f} B params, FIXED (no EMA, no grads)")
    # Verify vocabulary alignment via the tokenizer (Qwen3.5 multimodal
    # config nests vocab_size under text_config, so .config.vocab_size fails).
    s_vocab = len(processor.tokenizer)
    t_proc = AutoProcessor.from_pretrained(args.teacher_model, **t_kw)
    t_vocab = len(t_proc.tokenizer)
    if s_vocab != t_vocab:
        raise SystemExit(
            f"Vocab mismatch: teacher={t_vocab} vs student={s_vocab}. "
            f"Cannot do token-level KL.")
    print(f"  vocab match: {s_vocab} tokens ✓")

    # ── Dataset ─────────────────────────────────────────────────────────
    print(f"Loading dataset: {args.data_dir}")
    train_ds = build_dataset(args.data_dir / "train.jsonl", processor,
                             max_samples=args.max_samples, family="qwen",
                             official_prompt=args.official_prompt)
    print(f"  train={len(train_ds)}")

    collator = OnPolicyCollator(processor, max_length=args.max_seq_length,
                                teacher_mode=args.teacher_mode,
                                official_prompt=args.official_prompt)

    # Two-model variant — teacher already loaded above. EMA tracks the student.
    print(f"  EMA alpha={args.ema_alpha}  (separate teacher model on GPU)")

    # Stop-ids for on-policy generate
    tok = processor.tokenizer
    stop_ids = []
    for name in ("<|im_end|>", "<|endoftext|>"):
        tid = tok.convert_tokens_to_ids(name)
        if tid is not None and tid >= 0: stop_ids.append(tid)
    if tok.eos_token_id is not None and tok.eos_token_id not in stop_ids:
        stop_ids.append(tok.eos_token_id)
    print(f"  stop_ids: {stop_ids}")

    # ── Training loop (minimal, full HF) ────────────────────────────────
    from torch.utils.data import DataLoader
    dl = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True,
                    collate_fn=collator, num_workers=0)

    total_steps = (len(dl) // args.grad_accum) * args.epochs
    warmup_steps = max(int(args.warmup_ratio * total_steps), 10)
    print(f"  total_steps={total_steps} warmup={warmup_steps}")

    optim = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=args.lr, betas=(0.9, 0.999), weight_decay=0.0,
    )
    from torch.optim.lr_scheduler import LambdaLR
    def lr_lambda(step):
        if step < warmup_steps:
            return step / max(warmup_steps, 1)
        # cosine decay
        progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
        import math
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
    sched = LambdaLR(optim, lr_lambda)

    args.output_dir.mkdir(parents=True, exist_ok=True)

    # ── Resume optimizer + scheduler + global_step from training_state.pt ──
    # When --subfolder points to a checkpoint subfolder, attempt to download
    # `training_state.pt` from the same HF repo+subfolder and restore. The
    # model weights are already loaded above via from_pretrained(subfolder=…).
    global_step = 0
    if args.subfolder:
        try:
            from huggingface_hub import hf_hub_download
            state_file = hf_hub_download(
                repo_id=str(args.model),
                filename="training_state.pt",
                subfolder=args.subfolder,
                token=os.environ.get("HF_TOKEN"),
            )
            ts = torch.load(state_file, map_location="cpu")
            optim.load_state_dict(ts["optimizer"])
            sched.load_state_dict(ts["scheduler"])
            global_step = int(ts["global_step"])
            print(f"  → resumed: global_step={global_step}, lr={sched.get_last_lr()[0]:.2e}")
        except Exception as e:
            print(f"  NOTE: no training_state.pt at {args.model}/{args.subfolder} ({type(e).__name__}); "
                  f"warm-start from weights only, fresh optimizer/scheduler.")

    accum = 0
    loss_sum = 0.0
    n_kept = 0    # how many micro-batches in this grad-accum window contributed grads
    n_skip = 0    # rolling totals across training, for logs
    n_kept_total = 0
    for epoch in range(args.epochs):
        for batch in dl:
            loss, was_kept = train_step(
                model, teacher, batch,
                kl_type=args.kl_type, kl_temperature=args.kl_temperature,
                jsd_beta=args.jsd_beta,
                max_new_tokens=args.on_policy_max_new_tokens,
                sample_temp=args.on_policy_temperature, stop_ids=stop_ids,
                tokenizer=processor.tokenizer,
            )
            # Always call backward — when was_kept=False, loss is a 0-tensor
            # with requires_grad, so backward is a no-op on params.
            (loss / args.grad_accum).backward()
            loss_sum += float(loss.detach()) if was_kept else 0.0
            accum += 1
            if was_kept:
                n_kept += 1
                n_kept_total += 1
            else:
                n_skip += 1

            if accum == args.grad_accum:
                # Only step if at least one micro-batch in this window kept.
                if n_kept > 0:
                    torch.nn.utils.clip_grad_norm_(
                        [p for p in model.parameters() if p.requires_grad],
                        args.max_grad_norm)
                    optim.step(); sched.step()
                else:
                    sched.step()  # advance LR schedule even if no update
                optim.zero_grad()
                global_step += 1
                accum_kept = n_kept
                n_kept = 0; accum = 0

                if args.max_steps is not None and global_step >= args.max_steps:
                    break

                if global_step % args.log_every == 0:
                    avg = loss_sum / max(accum_kept, 1)
                    keep_rate = 100.0 * n_kept_total / max(n_kept_total + n_skip, 1)
                    print(f"step {global_step}/{total_steps} "
                          f"loss_{args.kl_type}={avg:.4f} lr={sched.get_last_lr()[0]:.2e} "
                          f"kept={accum_kept}/{args.grad_accum} keep_rate={keep_rate:.1f}%",
                          flush=True)
                    if args.wandb:
                        try:
                            import wandb
                            wandb.log({f"loss_{args.kl_type}": avg,
                                       "lr": sched.get_last_lr()[0],
                                       "kept_in_step": accum_kept,
                                       "keep_rate": keep_rate},
                                      step=global_step)
                        except Exception: pass
                    loss_sum = 0.0

                if global_step % args.save_steps == 0:
                    ck = args.output_dir / f"checkpoint-{global_step}"
                    ck.mkdir(exist_ok=True)
                    model.save_pretrained(ck)
                    processor.save_pretrained(ck)
                    # Save full training state (optimizer + scheduler + step)
                    # so a preempted pod can resume cleanly via --subfolder.
                    torch.save({
                        "optimizer": optim.state_dict(),
                        "scheduler": sched.state_dict(),
                        "global_step": global_step,
                    }, ck / "training_state.pt")
                    print(f"  → saved {ck} (with training_state.pt)")
                    if args.hf_repo:
                        try:
                            from huggingface_hub import HfApi
                            api = HfApi(token=os.environ.get("HF_TOKEN"))
                            api.create_repo(args.hf_repo, exist_ok=True, repo_type="model")
                            api.upload_folder(
                                folder_path=str(ck),
                                repo_id=args.hf_repo,
                                path_in_repo=f"checkpoint-{global_step}",
                                commit_message=f"checkpoint-{global_step}",
                            )
                            print(f"  → pushed checkpoint-{global_step} to {args.hf_repo}/checkpoint-{global_step}")
                        except Exception as e:
                            print(f"  WARNING: HF upload failed for checkpoint-{global_step}: {e}")

        if args.max_steps is not None and global_step >= args.max_steps:
            break

    # Final save
    model.save_pretrained(args.output_dir)
    processor.save_pretrained(args.output_dir)
    print(f"Training done. Final model → {args.output_dir}")
    if args.hf_repo:
        try:
            from huggingface_hub import HfApi
            api = HfApi(token=os.environ.get("HF_TOKEN"))
            api.create_repo(args.hf_repo, exist_ok=True, repo_type="model")
            api.upload_folder(
                folder_path=str(args.output_dir),
                repo_id=args.hf_repo,
                path_in_repo="final",
                commit_message="final model",
            )
            print(f"  → pushed final model to {args.hf_repo}/final")
        except Exception as e:
            print(f"  WARNING: HF upload failed for final model: {e}")


if __name__ == "__main__":
    main()