#!/usr/bin/env python3
"""VLAlert v3 SFT — Qwen3-VL-4B-Instruct + LoRA.

Fine-tunes Qwen3-VL-4B-Instruct with LoRA to produce structured safety
reasoning and belief tokens from dashcam frames.

Data format (v6_sft_train.jsonl):
  Each record has assistant_v6 text with [Analysis] + [Safety Assessment]
  sections, where beliefs are wrapped in <|BELIEF|>...</|BELIEF|> tags
  followed by action tokens (<|SILENT|>, <|OBSERVE|>, <|ALERT|>).

Loss:
  L_total = L_causal_LM + lambda_emb * L_belief_embedding

  L_causal_LM: weighted CE with belief content tokens at 1.5x and action
  tokens at 2.0x weight.

  L_belief_embedding: cosine similarity loss between hidden states at
  <|BELIEF|> positions and pre-computed target embeddings (optional,
  activated from --belief_emb_start_epoch onward).

Run:
  python -m training.VLA.train_vlalert_sft_v3 \
      --train_jsonl data/cot_corpus_v3/v6_sft_train.jsonl \
      --val_jsonl   data/cot_corpus_v3/v6_sft_val.jsonl \
      --out_dir     checkpoints/vlalert_sft \
      --epochs 3 --batch_size 4 --grad_accum 4 --lr 5e-5
"""
from __future__ import annotations

import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))

# ── Conv3d -> Linear patch (MUST run before any Qwen3-VL import) ──────────
import torch
import torch.nn as nn
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionPatchEmbed

_PATCH_APPLIED: dict = {}


def _fast_patch_embed_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """Conv3d -> Linear lazy replacement (math-identical)."""
    target_dtype = self.proj.weight.dtype
    if isinstance(self.proj, nn.Conv3d):
        conv = self.proj
        out_dim = conv.out_channels
        in_dim = (conv.in_channels * conv.kernel_size[0]
                  * conv.kernel_size[1] * conv.kernel_size[2])
        w_flat = conv.weight.detach().reshape(out_dim, in_dim).contiguous()
        bias = conv.bias.detach().clone() if conv.bias is not None else None
        new_proj = nn.Linear(in_dim, out_dim, bias=bias is not None)
        new_proj.weight.data.copy_(w_flat)
        if bias is not None:
            new_proj.bias.data.copy_(bias)
        new_proj.to(device=conv.weight.device, dtype=conv.weight.dtype)
        self.proj = new_proj
        if id(self) not in _PATCH_APPLIED:
            _PATCH_APPLIED[id(self)] = True
            print(f"[fast_patch] Conv3d({in_dim}->{out_dim}) -> "
                  f"Linear({in_dim}->{out_dim})", flush=True)
    if hidden_states.dim() > 2 or hidden_states.shape[-1] != self.proj.in_features:
        hidden_states = hidden_states.reshape(-1, self.proj.in_features)
    return self.proj(hidden_states.to(dtype=target_dtype))


Qwen3VLVisionPatchEmbed.forward = _fast_patch_embed_forward
print("[fast_patch] Qwen3VLVisionPatchEmbed.forward patched.", flush=True)

# ── standard imports ──────────────────────────────────────────────────────
import argparse
import json
import math
import os
import re
from typing import Any, Dict, List, Optional, Tuple

import cv2
import numpy as np
from PIL import Image
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

from peft import LoraConfig, get_peft_model
from transformers import AutoProcessor, AutoModelForImageTextToText
from transformers.optimization import get_cosine_schedule_with_warmup

# ── special tokens ────────────────────────────────────────────────────────

BELIEF_OPEN = "<|BELIEF|>"
BELIEF_CLOSE = "</|BELIEF|>"
ACTION_ALERT = "<|ALERT|>"
ACTION_OBSERVE = "<|OBSERVE|>"
ACTION_SILENT = "<|SILENT|>"
ACTION_TOKENS = [ACTION_ALERT, ACTION_OBSERVE, ACTION_SILENT]
ALL_SPECIAL = [BELIEF_OPEN, BELIEF_CLOSE] + ACTION_TOKENS

# ── prompts ───────────────────────────────────────────────────────────────

SYSTEM_PROMPT_V3 = (
    "You are a driving-safety assistant. Given N dashcam frames "
    "(earliest → latest), first analyze the safety situation for each frame "
    "in [Analysis], then produce a structured [Safety Assessment] with "
    "per-frame belief summaries wrapped in <|BELIEF|>...</|BELIEF|> "
    "and action tokens."
)


def user_prompt_v3(n_frames: int) -> str:
    return f"Analyze these {n_frames} frames for driving safety."


# ── frame loading ─────────────────────────────────────────────────────────

def _resize_bgr(frame: np.ndarray, resize_short: int) -> Image.Image:
    h, w = frame.shape[:2]
    scale = resize_short / min(h, w)
    nh, nw = int(round(h * scale)), int(round(w * scale))
    frame = cv2.resize(frame, (nw, nh), interpolation=cv2.INTER_AREA)
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    return Image.fromarray(frame)


def load_frames(video_path: str, frame_indices: List[int],
                resize_short: int = 336) -> List[Image.Image]:
    """Load specific frames from mp4 or image directory."""
    p = Path(video_path)

    if p.is_dir():
        # Image directory — check for images/ subdirectory (DoTA convention)
        if (p / "images").is_dir():
            p = p / "images"
        exts = (".jpg", ".jpeg", ".png")
        files = sorted([f for f in p.iterdir() if f.suffix.lower() in exts])
        if not files:
            raise RuntimeError(f"no images in {p}")
        total = len(files)
        frames: List[Image.Image] = []
        for idx in frame_indices:
            idx_clipped = max(0, min(total - 1, int(idx)))
            img = cv2.imread(str(files[idx_clipped]))
            if img is None:
                img = np.zeros((resize_short, resize_short, 3), dtype=np.uint8)
            frames.append(_resize_bgr(img, resize_short))
        return frames

    # mp4/mkv video file
    cap = cv2.VideoCapture(str(p))
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if total <= 0:
        cap.release()
        raise RuntimeError(f"bad video: {p}")
    clipped = [max(0, min(total - 1, int(i))) for i in frame_indices]
    wanted_sorted = sorted(set(clipped))
    picked: dict = {}
    cur = 0
    ptr = 0
    while cap.isOpened() and ptr < len(wanted_sorted):
        ok, frame = cap.read()
        if not ok:
            break
        while ptr < len(wanted_sorted) and cur == wanted_sorted[ptr]:
            picked[cur] = frame
            ptr += 1
        cur += 1
    cap.release()
    frames = []
    fallback = next(iter(picked.values())) if picked else None
    for i in clipped:
        f = picked.get(i, fallback)
        if f is None:
            f = np.zeros((resize_short, resize_short, 3), dtype=np.uint8)
        frames.append(_resize_bgr(f, resize_short))
    return frames


# ── chat template builders ────────────────────────────────────────────────

def build_chat_v3(frames: List[Image.Image], n_frames: int,
                  assistant_text: Optional[str] = None):
    """Build Qwen3-VL chat messages for v3 SFT."""
    user_content = [{"type": "image", "image": img} for img in frames]
    user_content.append({"type": "text", "text": user_prompt_v3(n_frames)})
    msgs = [
        {"role": "system",
         "content": [{"type": "text", "text": SYSTEM_PROMPT_V3}]},
        {"role": "user", "content": user_content},
    ]
    if assistant_text is not None:
        msgs.append({"role": "assistant",
                      "content": [{"type": "text", "text": assistant_text}]})
    return msgs


# ── Dataset ───────────────────────────────────────────────────────────────

class VLAlertSFTDatasetV3(Dataset):
    """v3 SFT dataset reading v6_sft_*.jsonl records.

    Each record has variable n_frames (1 or 8), assistant_v6 text with
    [Analysis] + [Safety Assessment] sections, and per-frame beliefs.
    """

    def __init__(self,
                 jsonl_path: str,
                 processor,
                 resize_short: int = 336,
                 max_len: int = 4096):
        self.processor = processor
        self.resize_short = resize_short
        self.max_len = max_len

        # cache token IDs for action and belief tokens
        tok = processor.tokenizer
        self.action_ids = set()
        for t in ACTION_TOKENS:
            tid = tok.convert_tokens_to_ids(t)
            if tid is not None and tid != tok.unk_token_id:
                self.action_ids.add(tid)

        self.belief_open_id = tok.convert_tokens_to_ids(BELIEF_OPEN)
        self.belief_close_id = tok.convert_tokens_to_ids(BELIEF_CLOSE)

        self.records: List[Dict[str, Any]] = []
        n_skipped = 0
        with open(jsonl_path) as f:
            for ln in f:
                ln = ln.strip()
                if not ln:
                    continue
                try:
                    r = json.loads(ln)
                except json.JSONDecodeError:
                    n_skipped += 1
                    continue
                # validate required fields
                ok = (r.get("video_path")
                      and isinstance(r.get("frame_indices"), list)
                      and len(r["frame_indices"]) > 0
                      and isinstance(r.get("n_frames"), int)
                      and r["n_frames"] == len(r["frame_indices"])
                      and isinstance(r.get("assistant_v6"), str)
                      and len(r["assistant_v6"].strip()) > 0
                      and isinstance(r.get("beliefs_per_frame"), list)
                      and len(r["beliefs_per_frame"]) == r["n_frames"])
                if not ok:
                    n_skipped += 1
                    continue
                self.records.append(r)
        print(f"[VLAlertSFTDatasetV3] loaded {len(self.records)} records "
              f"(skipped {n_skipped}) from {jsonl_path}")

    def __len__(self):
        return len(self.records)

    def __getitem__(self, idx):
        rec = self.records[idx]
        n_frames = rec["n_frames"]
        assistant_text = rec["assistant_v6"]

        # load frames
        frames = load_frames(rec["video_path"], rec["frame_indices"],
                             resize_short=self.resize_short)

        # build full (with assistant) and prefix (without) chat
        full_msgs = build_chat_v3(frames, n_frames, assistant_text)
        prefix_msgs = build_chat_v3(frames, n_frames, None)

        proc = self.processor
        full_text = proc.apply_chat_template(full_msgs, tokenize=False,
                                             add_generation_prompt=False)
        prefix_text = proc.apply_chat_template(prefix_msgs, tokenize=False,
                                               add_generation_prompt=True)

        full = proc(text=[full_text], images=[frames], return_tensors="pt",
                    padding=False, truncation=True, max_length=self.max_len)
        prefix = proc(text=[prefix_text], images=[frames], return_tensors="pt",
                      padding=False, truncation=True, max_length=self.max_len)

        input_ids = full["input_ids"][0]
        labels = input_ids.clone()
        prefix_len = prefix["input_ids"].shape[1]
        labels[:prefix_len] = -100

        # Build per-token weight masks
        # action_token_mask: True at action token positions
        # belief_content_mask: True at tokens BETWEEN <|BELIEF|> and </|BELIEF|>
        # (excluding the tags themselves)
        seq_len = input_ids.size(0)
        action_mask = torch.zeros(seq_len, dtype=torch.bool)
        belief_content_mask = torch.zeros(seq_len, dtype=torch.bool)
        belief_open_mask = torch.zeros(seq_len, dtype=torch.bool)

        ids_list = input_ids.tolist()
        in_belief = False
        for i in range(prefix_len, seq_len):
            tid = ids_list[i]
            if tid in self.action_ids:
                action_mask[i] = True
            if tid == self.belief_open_id:
                in_belief = True
                belief_open_mask[i] = True
                continue
            if tid == self.belief_close_id:
                in_belief = False
                continue
            if in_belief:
                belief_content_mask[i] = True

        item = {
            "input_ids": input_ids,
            "labels": labels,
            "action_token_mask": action_mask,
            "belief_content_mask": belief_content_mask,
            "belief_open_mask": belief_open_mask,
            "attention_mask": (full["attention_mask"][0]
                               if "attention_mask" in full else None),
            "pixel_values": (full["pixel_values"]
                             if "pixel_values" in full else None),
            "image_grid_thw": (full["image_grid_thw"]
                               if "image_grid_thw" in full else None),
            "record_idx": idx,
            "belief_source": rec.get("belief_source", ""),
            "n_beliefs": sum(1 for x in ids_list if x == self.belief_open_id),
        }
        for k in ("video_grid_thw", "pixel_values_videos"):
            if k in full:
                item[k] = full[k]
        return item


# ── Collator ──────────────────────────────────────────────────────────────

class CollatorV3:
    """Pad sequence dim; cat pixel/grid along their natural dim."""

    def __init__(self, processor):
        self.pad_id = (processor.tokenizer.pad_token_id
                       or processor.tokenizer.eos_token_id or 0)

    def __call__(self, batch):
        max_len = max(b["input_ids"].size(0) for b in batch)
        B = len(batch)
        ids = torch.full((B, max_len), self.pad_id, dtype=torch.long)
        labs = torch.full((B, max_len), -100, dtype=torch.long)
        action_mask = torch.zeros((B, max_len), dtype=torch.bool)
        belief_content_mask = torch.zeros((B, max_len), dtype=torch.bool)
        belief_open_mask = torch.zeros((B, max_len), dtype=torch.bool)
        attn_mask = torch.zeros((B, max_len), dtype=torch.long)

        record_idxs = []
        belief_sources = []
        n_beliefs_list = []

        for i, b in enumerate(batch):
            L = b["input_ids"].size(0)
            ids[i, :L] = b["input_ids"]
            labs[i, :L] = b["labels"]
            action_mask[i, :L] = b["action_token_mask"]
            belief_content_mask[i, :L] = b["belief_content_mask"]
            belief_open_mask[i, :L] = b["belief_open_mask"]
            if b.get("attention_mask") is not None:
                attn_mask[i, :L] = b["attention_mask"]
            else:
                attn_mask[i, :L] = 1
            record_idxs.append(b["record_idx"])
            belief_sources.append(b["belief_source"])
            n_beliefs_list.append(b["n_beliefs"])

        out = {
            "input_ids": ids,
            "labels": labs,
            "attention_mask": attn_mask,
            "action_token_mask": action_mask,
            "belief_content_mask": belief_content_mask,
            "belief_open_mask": belief_open_mask,
            "record_idxs": record_idxs,
            "belief_sources": belief_sources,
            "n_beliefs": n_beliefs_list,
        }
        # pixel_values: [num_patches_total, dim] -- cat across batch
        if batch[0].get("pixel_values") is not None:
            out["pixel_values"] = torch.cat(
                [b["pixel_values"] for b in batch], dim=0)
        # image_grid_thw: [n_images_per_sample, 3] -- cat across batch
        if batch[0].get("image_grid_thw") is not None:
            out["image_grid_thw"] = torch.cat(
                [b["image_grid_thw"] for b in batch], dim=0)
        for k in ("video_grid_thw", "pixel_values_videos"):
            if batch[0].get(k) is not None:
                out[k] = torch.cat([b[k] for b in batch], dim=0)
        return out


# ── Loss functions ────────────────────────────────────────────────────────

def weighted_ce_loss(logits: torch.Tensor,
                     labels: torch.Tensor,
                     action_mask: torch.Tensor,
                     belief_content_mask: torch.Tensor,
                     action_weight: float = 2.0,
                     belief_weight: float = 1.5) -> torch.Tensor:
    """Causal-LM CE with per-token weighting.

    Applies the standard next-token shift: position t's logits predict
    position (t+1)'s label.

    Weights:
      - Reasoning tokens ([Analysis] section): 1.0 (default)
      - Belief content tokens (between BELIEF tags): belief_weight
      - Action tokens (SILENT/OBSERVE/ALERT): action_weight
    """
    # shift for next-token prediction
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    shift_amask = action_mask[..., 1:].contiguous()
    shift_bmask = belief_content_mask[..., 1:].contiguous()

    V = shift_logits.size(-1)
    flat_logits = shift_logits.view(-1, V)
    flat_labels = shift_labels.view(-1)
    flat_amask = shift_amask.view(-1)
    flat_bmask = shift_bmask.view(-1)

    valid = flat_labels != -100
    if not valid.any():
        return flat_logits.sum() * 0.0

    loss_per = torch.nn.functional.cross_entropy(
        flat_logits[valid], flat_labels[valid], reduction="none")

    # build weight vector: start at 1.0, override for belief content and action
    w = torch.ones_like(loss_per)
    valid_bmask = flat_bmask[valid]
    valid_amask = flat_amask[valid]
    w = torch.where(valid_bmask,
                    torch.full_like(w, belief_weight), w)
    # action tokens override belief weight if both are set (they shouldn't overlap,
    # but action takes priority)
    w = torch.where(valid_amask,
                    torch.full_like(w, action_weight), w)

    return (loss_per * w).sum() / w.sum()


def belief_embedding_loss(hidden_states: torch.Tensor,
                          belief_open_mask: torch.Tensor,
                          record_idxs: List[int],
                          belief_sources: List[str],
                          n_beliefs: List[int],
                          belief_targets: torch.Tensor,
                          layer_idx: int = 28) -> torch.Tensor:
    """Cosine similarity loss on hidden states at <|BELIEF|> positions.

    Args:
        hidden_states: tuple of layer outputs from model (when output_hidden_states=True),
                       or a single tensor for the specified layer.
                       If tuple, index with layer_idx.
        belief_open_mask: [B, T] bool mask, True at <|BELIEF|> token positions.
        record_idxs: list of record indices in the dataset (length B).
        belief_sources: list of belief_source strings (length B).
        n_beliefs: list of number of beliefs per sample (length B).
        belief_targets: [N_records, max_beliefs, hidden_dim] pre-computed embeddings.
        layer_idx: which transformer layer's hidden state to use.

    Returns:
        Scalar loss = mean(1 - cos_sim(h_belief, target)) over valid positions.
    """
    # extract the right layer's hidden states
    if isinstance(hidden_states, (tuple, list)):
        h = hidden_states[layer_idx]  # [B, T, D]
    else:
        h = hidden_states  # already the right layer

    B, T, D = h.shape
    device = h.device

    total_loss = torch.tensor(0.0, device=device, dtype=h.dtype)
    count = 0

    for b_idx in range(B):
        src = belief_sources[b_idx]
        # only apply on records with gpt or annotation belief sources
        if "gpt" not in src and "annotation" not in src:
            continue

        rec_idx = record_idxs[b_idx]
        if rec_idx >= belief_targets.shape[0]:
            continue

        # find positions of <|BELIEF|> tokens in this sample
        positions = belief_open_mask[b_idx].nonzero(as_tuple=False).squeeze(-1)
        if positions.dim() == 0:
            positions = positions.unsqueeze(0)
        if positions.numel() == 0:
            continue

        n_b = min(positions.size(0), belief_targets.shape[1])
        for k in range(n_b):
            pos = positions[k].item()
            h_belief = h[b_idx, pos, :]  # [D]
            target = belief_targets[rec_idx, k, :].to(device=device, dtype=h.dtype)
            # skip zero targets (padding)
            if target.abs().sum() < 1e-8:
                continue
            cos_sim = torch.nn.functional.cosine_similarity(
                h_belief.unsqueeze(0), target.unsqueeze(0))
            total_loss = total_loss + (1.0 - cos_sim.squeeze())
            count += 1

    if count == 0:
        return torch.tensor(0.0, device=device, dtype=h.dtype, requires_grad=True)
    return total_loss / count


# ── CLI args ──────────────────────────────────────────────────────────────

def parse_args():
    ap = argparse.ArgumentParser(
        description="VLAlert v3 SFT: Qwen3-VL-4B + LoRA")
    ap.add_argument("--train_jsonl",
                    default="data/cot_corpus_v3/v6_sft_train.jsonl")
    ap.add_argument("--val_jsonl",
                    default="data/cot_corpus_v3/v6_sft_val.jsonl")
    ap.add_argument("--base_model",
                    default="models/Qwen3-VL-4B-Instruct")
    ap.add_argument("--out_dir",
                    default="checkpoints/vlalert_sft")
    ap.add_argument("--lora_r", type=int, default=64)
    ap.add_argument("--lora_alpha", type=int, default=128)
    ap.add_argument("--lora_dropout", type=float, default=0.05)
    ap.add_argument("--lr", type=float, default=5e-5)
    ap.add_argument("--epochs", type=int, default=3)
    ap.add_argument("--batch_size", type=int, default=4)
    ap.add_argument("--grad_accum", type=int, default=4)
    ap.add_argument("--warmup_ratio", type=float, default=0.03)
    ap.add_argument("--max_len", type=int, default=4096)
    ap.add_argument("--resize_short", type=int, default=336)
    ap.add_argument("--action_token_weight", type=float, default=2.0)
    ap.add_argument("--belief_token_weight", type=float, default=1.5)
    ap.add_argument("--belief_emb_weight", type=float, default=0.1,
                    help="Lambda for belief embedding loss")
    ap.add_argument("--belief_emb_start_epoch", type=int, default=2,
                    help="Epoch at which belief embedding loss activates "
                         "(1-indexed; default 2 = start at epoch 2)")
    ap.add_argument("--belief_target_path", type=str, default="",
                    help="Path to .pt file with pre-computed belief "
                         "target embeddings [N, max_beliefs, 2560]")
    ap.add_argument("--belief_layer_idx", type=int, default=28,
                    help="Transformer layer index for belief embedding extraction")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--max_samples", type=int, default=0,
                    help="Cap dataset size for smoke tests (0 = all)")
    ap.add_argument("--log_every", type=int, default=50)
    ap.add_argument("--save_every_epoch", action="store_true")
    ap.add_argument("--resume", type=str, default="",
                    help="Resume LoRA adapter from this directory")
    ap.add_argument("--num_workers", type=int, default=0)
    return ap.parse_args()


# ── token setup ───────────────────────────────────────────────────────────

def add_special_tokens_and_resize(processor, model):
    """Add BELIEF/action special tokens and resize embeddings."""
    tok = processor.tokenizer
    before = len(tok)
    added = tok.add_special_tokens({"additional_special_tokens": ALL_SPECIAL})
    after = len(tok)
    print(f"[tokens] vocab {before} -> {after}  ({added} new)")
    if added == 0:
        return
    model.resize_token_embeddings(after)
    emb = model.get_input_embeddings()
    with torch.no_grad():
        mean_vec = emb.weight[:before].mean(dim=0)
        for tok_str in ALL_SPECIAL:
            tid = tok.convert_tokens_to_ids(tok_str)
            emb.weight[tid] = mean_vec + 0.01 * torch.randn_like(mean_vec)
    out_emb = model.get_output_embeddings()
    if out_emb is not None and out_emb.weight.data_ptr() != emb.weight.data_ptr():
        with torch.no_grad():
            mean_out = out_emb.weight[:before].mean(dim=0)
            for tok_str in ALL_SPECIAL:
                tid = tok.convert_tokens_to_ids(tok_str)
                out_emb.weight[tid] = mean_out + 0.01 * torch.randn_like(mean_out)


# ── Validation ────────────────────────────────────────────────────────────

@torch.no_grad()
def validate(model, val_dl, args, epoch: int,
             belief_targets: Optional[torch.Tensor] = None) -> float:
    """Run one validation pass and return average loss."""
    model.eval()
    total_loss = 0.0
    total_steps = 0

    use_emb = (belief_targets is not None
               and args.belief_emb_weight > 0
               and (epoch + 1) >= args.belief_emb_start_epoch)

    for batch in tqdm(val_dl, desc=f"  val ep{epoch}", ncols=80, leave=False):
        input_ids = batch["input_ids"].to("cuda", non_blocking=True)
        labels = batch["labels"].to("cuda", non_blocking=True)
        amask = batch["action_token_mask"].to("cuda", non_blocking=True)
        bmask = batch["belief_content_mask"].to("cuda", non_blocking=True)
        bo_mask = batch["belief_open_mask"].to("cuda", non_blocking=True)
        attn = batch.get("attention_mask")
        if attn is not None:
            attn = attn.to("cuda", non_blocking=True)
        pix = batch["pixel_values"].to("cuda", dtype=torch.bfloat16,
                                       non_blocking=True)
        grid = batch["image_grid_thw"].to("cuda", non_blocking=True)

        fwd_kwargs = dict(input_ids=input_ids,
                          pixel_values=pix, image_grid_thw=grid)
        if attn is not None:
            fwd_kwargs["attention_mask"] = attn
        if use_emb:
            fwd_kwargs["output_hidden_states"] = True

        out = model(**fwd_kwargs)

        ce_loss = weighted_ce_loss(
            out.logits, labels, amask, bmask,
            action_weight=args.action_token_weight,
            belief_weight=args.belief_token_weight)

        loss = ce_loss

        if use_emb and hasattr(out, "hidden_states") and out.hidden_states is not None:
            emb_loss = belief_embedding_loss(
                out.hidden_states, bo_mask,
                batch["record_idxs"], batch["belief_sources"],
                batch["n_beliefs"], belief_targets,
                layer_idx=args.belief_layer_idx)
            loss = loss + args.belief_emb_weight * emb_loss

        total_loss += loss.float().item()
        total_steps += 1

    model.train()
    avg = total_loss / max(1, total_steps)
    return avg


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

def main():
    args = parse_args()
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    # resolve paths relative to ROOT
    base_model = args.base_model
    if not os.path.isabs(base_model):
        base_model = str(ROOT / base_model)

    out_dir = Path(args.out_dir)
    if not out_dir.is_absolute():
        out_dir = ROOT / out_dir
    out_dir.mkdir(parents=True, exist_ok=True)

    train_jsonl = args.train_jsonl
    if not os.path.isabs(train_jsonl):
        train_jsonl = str(ROOT / train_jsonl)
    val_jsonl = args.val_jsonl
    if not os.path.isabs(val_jsonl):
        val_jsonl = str(ROOT / val_jsonl)

    print(f"[train] base_model = {base_model}")
    print(f"[train] out_dir    = {out_dir}")
    print(f"[train] train_jsonl = {train_jsonl}")
    print(f"[train] val_jsonl   = {val_jsonl}")

    # ── load processor + model ────────────────────────────────────────
    print(f"[train] loading processor/model from {base_model}")
    processor = AutoProcessor.from_pretrained(base_model, trust_remote_code=True)
    model = AutoModelForImageTextToText.from_pretrained(
        base_model, torch_dtype=torch.bfloat16,
        trust_remote_code=True, attn_implementation="sdpa",
    )

    add_special_tokens_and_resize(processor, model)

    # freeze vision tower
    for attr in ("visual", "vision_tower"):
        if hasattr(model, attr):
            for p in getattr(model, attr).parameters():
                p.requires_grad = False

    if hasattr(model, "enable_input_require_grads"):
        model.enable_input_require_grads()
    model.gradient_checkpointing_enable(
        gradient_checkpointing_kwargs={"use_reentrant": False})

    # ── LoRA / resume ─────────────────────────────────────────────────
    if args.resume:
        from peft import PeftModel
        print(f"[resume] loading PEFT adapter from {args.resume}")
        model = PeftModel.from_pretrained(model, args.resume, is_trainable=True)
    else:
        lora_cfg = LoraConfig(
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            lora_dropout=args.lora_dropout,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                            "gate_proj", "up_proj", "down_proj"],
            bias="none",
            task_type="CAUSAL_LM",
            modules_to_save=["embed_tokens", "lm_head"],
        )
        model = get_peft_model(model, lora_cfg)
    model.print_trainable_parameters()
    model.to("cuda")
    model.config.use_cache = False

    # ── Load belief target embeddings (optional) ──────────────────────
    belief_targets = None
    if args.belief_target_path:
        bt_path = args.belief_target_path
        if not os.path.isabs(bt_path):
            bt_path = str(ROOT / bt_path)
        if os.path.exists(bt_path):
            belief_targets = torch.load(bt_path, map_location="cpu",
                                        weights_only=True)
            print(f"[train] loaded belief targets: {belief_targets.shape} "
                  f"from {bt_path}")
        else:
            print(f"[warn] belief_target_path not found: {bt_path}, "
                  f"embedding loss disabled")

    # ── Datasets ──────────────────────────────────────────────────────
    train_ds = VLAlertSFTDatasetV3(
        jsonl_path=train_jsonl, processor=processor,
        resize_short=args.resize_short, max_len=args.max_len)

    if args.max_samples > 0 and len(train_ds) > args.max_samples:
        from torch.utils.data import Subset
        train_ds = Subset(train_ds, list(range(args.max_samples)))
        print(f"[smoke] truncated training to {len(train_ds)} samples")

    val_ds = VLAlertSFTDatasetV3(
        jsonl_path=val_jsonl, processor=processor,
        resize_short=args.resize_short, max_len=args.max_len)

    print(f"[train] train={len(train_ds)}, val={len(val_ds)}")

    collator = CollatorV3(processor)
    train_dl = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True,
                          num_workers=args.num_workers, collate_fn=collator,
                          pin_memory=True)
    val_dl = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False,
                        num_workers=args.num_workers, collate_fn=collator,
                        pin_memory=True)

    # ── Optimizer + scheduler ─────────────────────────────────────────
    trainable = [p for p in model.parameters() if p.requires_grad]
    opt = AdamW(trainable, lr=args.lr, betas=(0.9, 0.95), weight_decay=0.0)
    total_updates = math.ceil(len(train_dl) * args.epochs / args.grad_accum)
    warmup = max(1, int(total_updates * args.warmup_ratio))
    sched = get_cosine_schedule_with_warmup(opt, warmup, total_updates)
    print(f"[train] total_updates={total_updates} warmup={warmup} lr={args.lr}")
    print(f"[train] belief_emb_weight={args.belief_emb_weight} "
          f"start_epoch={args.belief_emb_start_epoch}")

    # ── Training loop ─────────────────────────────────────────────────
    best_val_loss = float("inf")
    global_step = 0
    model.train()

    for epoch in range(args.epochs):
        # determine whether embedding loss is active this epoch
        use_emb = (belief_targets is not None
                   and args.belief_emb_weight > 0
                   and (epoch + 1) >= args.belief_emb_start_epoch)

        pbar = tqdm(enumerate(train_dl), total=len(train_dl),
                    desc=f"ep{epoch}", ncols=100, leave=True)
        running_ce = 0.0
        running_emb = 0.0
        running_n = 0

        for step, batch in pbar:
            input_ids = batch["input_ids"].to("cuda", non_blocking=True)
            labels = batch["labels"].to("cuda", non_blocking=True)
            amask = batch["action_token_mask"].to("cuda", non_blocking=True)
            bmask = batch["belief_content_mask"].to("cuda", non_blocking=True)
            bo_mask = batch["belief_open_mask"].to("cuda", non_blocking=True)
            attn = batch.get("attention_mask")
            if attn is not None:
                attn = attn.to("cuda", non_blocking=True)
            pix = batch["pixel_values"].to("cuda", dtype=torch.bfloat16,
                                           non_blocking=True)
            grid = batch["image_grid_thw"].to("cuda", non_blocking=True)

            fwd_kwargs = dict(input_ids=input_ids,
                              pixel_values=pix, image_grid_thw=grid)
            if attn is not None:
                fwd_kwargs["attention_mask"] = attn
            if use_emb:
                fwd_kwargs["output_hidden_states"] = True

            out = model(**fwd_kwargs)

            # L_causal_LM
            ce_loss = weighted_ce_loss(
                out.logits, labels, amask, bmask,
                action_weight=args.action_token_weight,
                belief_weight=args.belief_token_weight)

            loss = ce_loss

            # L_belief_embedding (optional)
            emb_loss_val = 0.0
            if use_emb and hasattr(out, "hidden_states") and out.hidden_states is not None:
                emb_loss = belief_embedding_loss(
                    out.hidden_states, bo_mask,
                    batch["record_idxs"], batch["belief_sources"],
                    batch["n_beliefs"], belief_targets,
                    layer_idx=args.belief_layer_idx)
                loss = loss + args.belief_emb_weight * emb_loss
                emb_loss_val = emb_loss.detach().float().item()

            loss = loss / args.grad_accum
            loss.backward()

            running_ce += ce_loss.detach().float().item()
            running_emb += emb_loss_val
            running_n += 1

            if (step + 1) % args.grad_accum == 0 or (step + 1) == len(train_dl):
                torch.nn.utils.clip_grad_norm_(trainable, 1.0)
                opt.step()
                sched.step()
                opt.zero_grad(set_to_none=True)
                global_step += 1

                if global_step % args.log_every == 0:
                    avg_ce = running_ce / max(1, running_n)
                    avg_emb = running_emb / max(1, running_n)
                    lr_now = sched.get_last_lr()[0]
                    pbar.set_postfix(
                        ce=f"{avg_ce:.4f}",
                        emb=f"{avg_emb:.4f}" if use_emb else "off",
                        lr=f"{lr_now:.2e}",
                        step=global_step)
                    running_ce, running_emb, running_n = 0.0, 0.0, 0

        # ── save checkpoint BEFORE validation (防止validation崩溃丢失训练) ──
        if args.save_every_epoch:
            ep_dir = out_dir / f"epoch_{epoch}"
            ep_dir.mkdir(parents=True, exist_ok=True)
            model.save_pretrained(ep_dir)
            processor.save_pretrained(ep_dir)
            with (ep_dir / "train_args.json").open("w") as f:
                json.dump(vars(args), f, indent=2)
            print(f"[save] epoch {epoch} -> {ep_dir} (pre-validation)")

        # ── end-of-epoch validation ───────────────────────────────────
        try:
            val_loss = validate(model, val_dl, args, epoch, belief_targets)
        except Exception as e:
            print(f"[warn] validation failed: {e}")
            val_loss = float("inf")
        print(f"[val] epoch {epoch}: val_loss={val_loss:.4f} "
              f"(best={best_val_loss:.4f})")

        # save best
        is_best = val_loss < best_val_loss
        if is_best:
            best_val_loss = val_loss
            best_dir = out_dir / "best"
            best_dir.mkdir(parents=True, exist_ok=True)
            model.save_pretrained(best_dir)
            processor.save_pretrained(best_dir)
            with (best_dir / "train_args.json").open("w") as f:
                json.dump(vars(args), f, indent=2)
            with (best_dir / "belief_tokens.json").open("w") as f:
                json.dump({
                    "special_tokens": ALL_SPECIAL,
                    "belief_open": BELIEF_OPEN,
                    "belief_close": BELIEF_CLOSE,
                    "actions": ACTION_TOKENS,
                }, f, indent=2)
            print(f"[save] best -> {best_dir} (val_loss={val_loss:.4f})")

        # (epoch checkpoint already saved pre-validation)
        if False:  # disabled — already saved above
            with (ep_dir / "train_args.json").open("w") as f:
                json.dump(vars(args), f, indent=2)
            print(f"[save] epoch {epoch} -> {ep_dir}")

    # ── Final save (last epoch) ───────────────────────────────────────
    final_dir = out_dir / "final"
    final_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(final_dir)
    processor.save_pretrained(final_dir)
    with (final_dir / "train_args.json").open("w") as f:
        json.dump(vars(args), f, indent=2)
    with (final_dir / "belief_tokens.json").open("w") as f:
        json.dump({
            "special_tokens": ALL_SPECIAL,
            "belief_open": BELIEF_OPEN,
            "belief_close": BELIEF_CLOSE,
            "actions": ACTION_TOKENS,
        }, f, indent=2)
    print(f"[done] final -> {final_dir}")
    print(f"[done] best val_loss = {best_val_loss:.4f}")


if __name__ == "__main__":
    main()