training/SFT/trainer.py

"""
SFT Trainer for TTA (Time-to-Accident) Regression
- Loads Qwen2.5-VL backbone + LoRA
- Trains belief_aggregator + TTA head
- Supports resuming from SFT checkpoints (LoRA + heads + optional optimizer state)
- Robust LoRA grad/update checks (no false-positive with grad accumulation / bf16 tiny updates)

NEW in this version (for your request):
1) Reset best_val_loss when resuming (default: ON)
2) Optionally run a fresh evaluation on the NEW val dataset immediately after resume (default: ON)
   so "best" is re-defined under the new val split.
"""

from __future__ import annotations

import os
import json
import time
import math
import random
import logging
import argparse
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any
from collections import defaultdict

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.amp import GradScaler, autocast
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
from tqdm import tqdm

# Optional deps
try:
    import wandb
    HAS_WANDB = True
except Exception:
    HAS_WANDB = False
    wandb = None

try:
    from transformers import AutoProcessor, AutoModelForVision2Seq
    from peft import PeftModel, LoraConfig, get_peft_model
    HAS_TRANSFORMERS = True
except Exception:
    HAS_TRANSFORMERS = False

# Local imports
from .dataset import SFTDataset, sft_collate_fn


# ---------------- Logging ----------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("SFT.trainer")


# ============================================================================
# Model Components
# ============================================================================

class HazardHead(nn.Module):
    """Binary hazard head: outputs hazard_prob ∈ (0, 1).

    Initialized to be slightly below 0.5 (lean toward safe at start).
    """

    def __init__(self, hidden_dim: int):
        super().__init__()
        self.fc = nn.Linear(hidden_dim, 1)
        nn.init.zeros_(self.fc.weight)
        self.fc.bias.data = torch.tensor([-1.0])   # sigmoid(-1) ≈ 0.27

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        """Returns hazard_logit [B] (raw, pre-sigmoid)."""
        return self.fc(hidden_state).squeeze(-1)


class TTAHead(nn.Module):
    """TTA Regression Head: outputs (tta_mean, tta_logvar)."""

    def __init__(self, hidden_dim: int, intermediate_dim: int = 512, dropout: float = 0.1):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.intermediate_dim = intermediate_dim

        self.net = nn.Sequential(
            nn.Linear(hidden_dim, intermediate_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(intermediate_dim, intermediate_dim // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(intermediate_dim // 2, 2),
        )

        self._init_weights()

    def _init_weights(self):
        nn.init.zeros_(self.net[-1].weight)
        # bias: mean=5, logvar=0
        self.net[-1].bias.data = torch.tensor([5.0, 0.0])

    def forward(self, hidden_state: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        out = self.net(hidden_state)
        tta_mean = F.softplus(out[:, 0])
        tta_logvar = out[:, 1]
        return tta_mean, tta_logvar


class BeliefAggregator(nn.Module):
    """Aggregate token hidden states to a single belief vector.

    Strategies:
      - mean_pool      : masked mean over all tokens  -> [B, D]
      - last_token     : hidden at last real token    -> [B, D]
      - attention_pool : learned-query attention pool -> [B, D]
      - dual_pool      : [mean(image_tokens) || mean(text_tokens)] -> [B, 2D]
                         Requires image_token_id (and optionally video_token_id).
                         This is P0.2 L1 "dual-modality pooling" — prevents the
                         language prompt from being diluted 10× by image tokens.
    """

    def __init__(
        self,
        hidden_dim: int,
        strategy: str = "mean_pool",
        image_token_id: Optional[int] = None,
        video_token_id: Optional[int] = None,
    ):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.strategy = strategy
        self.image_token_id = image_token_id
        self.video_token_id = video_token_id

        if strategy == "attention_pool":
            self.query = nn.Parameter(torch.randn(1, 1, hidden_dim) * 0.02)
            self.key_proj = nn.Linear(hidden_dim, hidden_dim)

        if strategy == "dual_pool" and image_token_id is None and video_token_id is None:
            raise ValueError("dual_pool requires image_token_id and/or video_token_id.")

    @property
    def output_dim(self) -> int:
        return 2 * self.hidden_dim if self.strategy == "dual_pool" else self.hidden_dim

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        input_ids: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if self.strategy == "mean_pool":
            return self._mean_pool(hidden_states, attention_mask)
        if self.strategy == "last_token":
            return self._last_token(hidden_states, attention_mask)
        if self.strategy == "attention_pool":
            return self._attention_pool(hidden_states, attention_mask)
        if self.strategy == "dual_pool":
            if input_ids is None:
                raise RuntimeError("dual_pool requires input_ids to separate image vs text tokens.")
            return self._dual_pool(hidden_states, attention_mask, input_ids)
        raise ValueError(f"Unknown strategy: {self.strategy}")

    def _mean_pool(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]) -> torch.Tensor:
        if attention_mask is None:
            return hidden_states.mean(dim=1)
        mask = attention_mask.unsqueeze(-1).float()
        masked = hidden_states * mask
        denom = mask.sum(dim=1).clamp(min=1e-9)
        return masked.sum(dim=1) / denom

    def _last_token(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]) -> torch.Tensor:
        if attention_mask is None:
            return hidden_states[:, -1, :]
        seq_lens = attention_mask.sum(dim=1).long() - 1
        b = torch.arange(hidden_states.size(0), device=hidden_states.device)
        return hidden_states[b, seq_lens, :]

    def _attention_pool(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]) -> torch.Tensor:
        B, L, D = hidden_states.shape
        q = self.query.expand(B, -1, -1)
        k = self.key_proj(hidden_states)
        scores = torch.bmm(q, k.transpose(1, 2)) / math.sqrt(D)
        if attention_mask is not None:
            scores = scores.masked_fill(attention_mask.unsqueeze(1) == 0, -1e9)
        w = F.softmax(scores, dim=-1)
        return torch.bmm(w, hidden_states).squeeze(1)

    def _dual_pool(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor],
        input_ids: torch.Tensor,
    ) -> torch.Tensor:
        """Separately mean-pool image tokens and text tokens, concat -> [B, 2D]."""
        is_img = torch.zeros_like(input_ids, dtype=torch.bool)
        if self.image_token_id is not None:
            is_img = is_img | (input_ids == self.image_token_id)
        if self.video_token_id is not None:
            is_img = is_img | (input_ids == self.video_token_id)

        if attention_mask is not None:
            valid = attention_mask > 0
            is_img = is_img & valid
            is_text = (~is_img) & valid
        else:
            is_text = ~is_img

        def _masked_mean(mask_bool: torch.Tensor) -> torch.Tensor:
            m = mask_bool.unsqueeze(-1).to(hidden_states.dtype)
            s = (hidden_states * m).sum(dim=1)
            denom = m.sum(dim=1).clamp(min=1e-6)
            return s / denom

        img_pool = _masked_mean(is_img)
        text_pool = _masked_mean(is_text)
        return torch.cat([img_pool, text_pool], dim=-1)


# ============================================================================
# SFT Model
# ============================================================================

class SFTModel(nn.Module):
    """VLM + LoRA + belief aggregator + HazardHead + TTAHead (dual head)."""

    def __init__(
        self,
        model_name: str = "Qwen/Qwen2.5-VL-3B-Instruct",
        pretrained_lora_path: Optional[str] = None,
        belief_strategy: str = "mean_pool",
        tta_intermediate_dim: int = 512,
        use_lora: bool = True,
        lora_r: int = 32,
        lora_alpha: int = 64,
        lora_dropout: float = 0.1,
        lora_target_modules: Optional[List[str]] = None,
        use_bf16: bool = True,
        device: str = "auto",
        max_pixels: Optional[int] = None,   # None → 768*28*28 default
        # P0.3 PEFT upgrade flags
        use_dora: bool = False,
        use_rslora: bool = False,
        lora_init: str = "default",     # "default" | "pissa" | "pissa_niter_16" | "olora" | "gaussian"
        attn_implementation: str = "flash_attention_2",
    ):
        super().__init__()
        if not HAS_TRANSFORMERS:
            raise RuntimeError("transformers/peft not available in this env.")

        self.model_name = model_name
        self.use_lora = use_lora
        self.use_bf16 = use_bf16

        if lora_target_modules is None:
            lora_target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]

        dtype = torch.bfloat16 if use_bf16 else torch.float32

        logger.info(f"📦 Loading VLM: {model_name}  (attn={attn_implementation})")
        self.vlm = AutoModelForVision2Seq.from_pretrained(
            model_name,
            torch_dtype=dtype,
            device_map="cuda:0",
            trust_remote_code=True,
            attn_implementation=attn_implementation,
        )

        if hasattr(self.vlm, "config"):
            self.vlm.config.use_cache = False

        if hasattr(self.vlm, "gradient_checkpointing_enable"):
            try:
                self.vlm.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
            except TypeError:
                self.vlm.gradient_checkpointing_enable()

        if hasattr(self.vlm, "enable_input_require_grads"):
            try:
                self.vlm.enable_input_require_grads()
            except Exception:
                pass

        _min_pixels = 256 * 28 * 28
        _max_pixels = max_pixels if max_pixels is not None else (768 * 28 * 28)
        logger.info(f"  max_pixels: {_max_pixels}  ({_max_pixels // (28*28)} tokens/frame max)")
        self.processor = AutoProcessor.from_pretrained(
            model_name,
            trust_remote_code=True,
            min_pixels=_min_pixels,
            max_pixels=_max_pixels,
        )

        self.hidden_dim = getattr(self.vlm.config, "hidden_size", None)
        if self.hidden_dim is None:
            raise RuntimeError("Cannot infer hidden_size from model config.")
        logger.info(f"  Hidden dim: {self.hidden_dim}")

        if use_lora:
            if pretrained_lora_path is not None:
                p = Path(pretrained_lora_path)
                if (p / "adapter_config.json").exists() and (p / "adapter_model.safetensors").exists():
                    logger.info(f"  Loading pretrained LoRA via PeftModel.from_pretrained: {p}")
                    self.vlm = PeftModel.from_pretrained(self.vlm, str(p), is_trainable=True)
                else:
                    logger.warning(f"⚠️ pretrained_lora_path exists but missing adapter files: {p}. Creating new LoRA.")
                    pretrained_lora_path = None

            if pretrained_lora_path is None:
                logger.info(
                    f"  Creating new LoRA (r={lora_r}, alpha={lora_alpha}, dropout={lora_dropout}, "
                    f"use_dora={use_dora}, use_rslora={use_rslora}, init={lora_init})"
                )
                lora_kwargs = dict(
                    r=lora_r,
                    lora_alpha=lora_alpha,
                    target_modules=lora_target_modules,
                    lora_dropout=lora_dropout,
                    bias="none",
                    task_type="CAUSAL_LM",
                )
                if use_dora:
                    lora_kwargs["use_dora"] = True
                if use_rslora:
                    lora_kwargs["use_rslora"] = True
                if lora_init and lora_init != "default":
                    # peft accepts: True | False | "gaussian" | "olora" | "pissa" | "pissa_niter_[N]" | "loftq"
                    lora_kwargs["init_lora_weights"] = lora_init
                lora_config = LoraConfig(**lora_kwargs)
                self.vlm = get_peft_model(self.vlm, lora_config)

            base = self.get_base_model()
            if hasattr(base, "config"):
                base.config.use_cache = False
            if hasattr(base, "gradient_checkpointing_enable"):
                try:
                    base.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
                except TypeError:
                    base.gradient_checkpointing_enable()
            if hasattr(base, "enable_input_require_grads"):
                try:
                    base.enable_input_require_grads()
                except Exception:
                    pass

            try:
                self.vlm.print_trainable_parameters()
            except Exception:
                pass

        self._register_requires_grad_hooks()

        self.device = next(self.vlm.parameters()).device
        self.dtype = next(self.vlm.parameters()).dtype

        # Grab image / video token ids from the VLM config (Qwen2.5-VL: 151655 / 151656).
        _cfg = getattr(self.vlm, "config", None)
        img_tok_id = getattr(_cfg, "image_token_id", None)
        vid_tok_id = getattr(_cfg, "video_token_id", None)
        if img_tok_id is None:
            img_tok_id = 151655  # Qwen2.5-VL fallback
        if vid_tok_id is None:
            vid_tok_id = 151656  # Qwen2.5-VL fallback

        self.belief_aggregator = BeliefAggregator(
            self.hidden_dim,
            strategy=belief_strategy,
            image_token_id=img_tok_id,
            video_token_id=vid_tok_id,
        ).to(self.device, dtype=self.dtype)

        belief_dim = self.belief_aggregator.output_dim
        self.belief_dim = belief_dim
        self.hazard_head = HazardHead(belief_dim).to(self.device, dtype=self.dtype)
        self.tta_head = TTAHead(belief_dim, intermediate_dim=tta_intermediate_dim).to(self.device, dtype=self.dtype)

        trainable = [(n, p) for n, p in self.vlm.named_parameters() if p.requires_grad]
        lora_trainable = [(n, p) for n, p in trainable if "lora_" in n.lower()]
        logger.info(f"  Trainable tensors: {len(trainable)}; LoRA trainable tensors: {len(lora_trainable)}")

        logger.info("✅ SFTModel initialized")
        logger.info(f"  Device: {self.device}")
        logger.info(f"  Dtype: {self.dtype}")
        logger.info(f"  Belief strategy: {belief_strategy}")

    def get_base_model(self):
        if hasattr(self.vlm, "get_base_model"):
            try:
                return self.vlm.get_base_model()
            except Exception:
                pass
        return getattr(self.vlm, "model", self.vlm)

    def _register_requires_grad_hooks(self):
        def _force_requires_grad_hook(_module, _inp, out):
            try:
                if torch.is_tensor(out) and out.is_floating_point():
                    out.requires_grad_(True)
                elif isinstance(out, (tuple, list)):
                    for t in out:
                        if torch.is_tensor(t) and t.is_floating_point():
                            t.requires_grad_(True)
            except Exception:
                return

        base_model = self.get_base_model()

        try:
            emb = base_model.get_input_embeddings() if hasattr(base_model, "get_input_embeddings") else None
            if emb is not None:
                emb.register_forward_hook(_force_requires_grad_hook)
                logger.info("✅ Registered requires_grad hook on TEXT embeddings")
        except Exception as e:
            logger.warning(f"⚠️ Failed to hook TEXT embeddings: {e}")

        try:
            hooked = False
            for name in ["visual", "vision_tower", "vision_model", "vision_encoder"]:
                if hasattr(base_model, name):
                    getattr(base_model, name).register_forward_hook(_force_requires_grad_hook)
                    logger.info(f"✅ Registered requires_grad hook on VISION module: {name}")
                    hooked = True
                    break
            if not hooked:
                for n, m in base_model.named_modules():
                    nl = n.lower()
                    if any(k in nl for k in ["visual", "vision", "patch_embed", "patch_embedding", "img_embed"]):
                        m.register_forward_hook(_force_requires_grad_hook)
                        logger.info(f"✅ Registered requires_grad hook on VISION submodule: {n}")
                        break
        except Exception as e:
            logger.warning(f"⚠️ Failed to hook VISION module: {e}")

    def encode_observation(self, batch_inputs: Dict[str, torch.Tensor]) -> torch.Tensor:
        moved: Dict[str, Any] = {}
        for k, v in batch_inputs.items():
            if not isinstance(v, torch.Tensor):
                moved[k] = v
                continue
            if k == "pixel_values":
                moved[k] = v.to(self.device, dtype=self.dtype, non_blocking=True)
            else:
                moved[k] = v.to(self.device, non_blocking=True)

        base = self.get_base_model()

        hidden_states = None
        core = getattr(base, "model", None)
        if core is not None:
            try:
                out = core(
                    input_ids=moved["input_ids"],
                    attention_mask=moved.get("attention_mask"),
                    pixel_values=moved.get("pixel_values"),
                    image_grid_thw=moved.get("image_grid_thw"),
                    use_cache=False,
                    return_dict=True,
                )
                hidden_states = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
            except TypeError:
                hidden_states = None

        if hidden_states is None:
            out = base(
                input_ids=moved["input_ids"],
                attention_mask=moved.get("attention_mask"),
                pixel_values=moved.get("pixel_values"),
                image_grid_thw=moved.get("image_grid_thw"),
                use_cache=False,
                return_dict=True,
                output_hidden_states=True,
            )
            if not hasattr(out, "hidden_states") or out.hidden_states is None:
                raise RuntimeError("Model output has no hidden_states; cannot build belief.")
            hidden_states = out.hidden_states[-1]

        if hidden_states.dim() != 3 or hidden_states.size(-1) != self.hidden_dim:
            raise RuntimeError(f"Unexpected hidden_states shape {tuple(hidden_states.shape)}, expected [B,L,{self.hidden_dim}]")

        return self.belief_aggregator(
            hidden_states,
            moved.get("attention_mask"),
            moved.get("input_ids"),
        )

    def forward(self, batch_inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        belief = self.encode_observation(batch_inputs)
        hazard_logit          = self.hazard_head(belief)          # raw logit [B]
        hazard_prob           = torch.sigmoid(hazard_logit)       # probability [B]
        tta_mean, tta_logvar  = self.tta_head(belief)
        return {
            "hazard_logit": hazard_logit,
            "hazard_prob":  hazard_prob,
            "tta_mean":     tta_mean,
            "tta_logvar":   tta_logvar,
            "belief":       belief.detach(),
        }

    def save_checkpoint(self, save_dir: str, epoch: int = 0, step: int = 0):
        save_dir = Path(save_dir)
        save_dir.mkdir(parents=True, exist_ok=True)

        if self.use_lora:
            lora_dir = save_dir / "vlm_lora"
            self.vlm.save_pretrained(lora_dir)
            logger.info(f"  Saved LoRA to {lora_dir}")

        torch.save(self.belief_aggregator.state_dict(), save_dir / "belief_aggregator.pt")
        torch.save(self.hazard_head.state_dict(),       save_dir / "hazard_head.pt")
        torch.save(self.tta_head.state_dict(),           save_dir / "tta_head.pt")

        cfg = {
            "model_name": self.model_name,
            "hidden_dim": self.hidden_dim,
            "belief_strategy": self.belief_aggregator.strategy,
            "belief_dim": self.belief_aggregator.output_dim,
            "image_token_id": self.belief_aggregator.image_token_id,
            "video_token_id": self.belief_aggregator.video_token_id,
            "tta_intermediate_dim": self.tta_head.intermediate_dim,
            "epoch": epoch,
            "step": step,
        }
        with open(save_dir / "config.json", "w") as f:
            json.dump(cfg, f, indent=2)

        logger.info(f"✅ Checkpoint saved to {save_dir}")


# ============================================================================
# Loss
# ============================================================================

def compute_sft_loss(
    hazard_logit: torch.Tensor,
    tta_mean: torch.Tensor,
    tta_logvar: torch.Tensor,
    hazard_label: torch.Tensor,
    hazard_weight: torch.Tensor,
    is_ego_positive: torch.Tensor,
    is_censored: torch.Tensor,
    tta_label: torch.Tensor,
    tta_cap: float = 10.0,
    nll_weight: float = 0.5,
    tta_obs_weight: float = 1.0,
    tta_cens_weight: float = 0.5,
    # legacy kwarg kept for callers that still pass hazard_prob
    hazard_prob: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, Dict[str, float]]:
    """
    Dual-head SFT loss.

    Hazard head (all samples):
        L_hazard = weighted_BCE_with_logits(hazard_logit, hazard_label)
        weights: ego_pos=1.0, safe_neg=1.0, non_ego=0.35, pre_risky=0.8

    TTA head (ego_positive only):
        Observed (TTA ≤ 10s): MSE + nll_weight * NLL
        Censored (TTA > 10s): relu(tta_cap - tta_mean)²

    Non-ego and safe_neg: NO TTA gradient.
    """
    hl   = hazard_label.float()
    hw   = hazard_weight.float()
    tm   = tta_mean.float()
    tlv  = tta_logvar.float()
    tl   = tta_label.float().clamp(0.1, tta_cap)
    var  = torch.exp(tlv).clamp(min=1e-6)
    zero = tta_mean.new_zeros(())

    # ── hazard loss (logits → safe with autocast) ─────────────────────────
    hl_logit = hazard_logit.float()
    bce_unreduced = F.binary_cross_entropy_with_logits(hl_logit, hl, reduction="none")
    L_hazard = (bce_unreduced * hw).mean()

    # keep hp for metrics
    hp = torch.sigmoid(hl_logit).detach()

    # ── TTA loss: ego_positive observed ──────────────────────────────────────
    obs_mask  = is_ego_positive & (~is_censored)
    cens_mask = is_ego_positive & is_censored

    if obs_mask.any():
        m   = tm[obs_mask];  l  = tl[obs_mask]
        lv  = tlv[obs_mask]; v  = var[obs_mask]
        mse = F.mse_loss(m, l)
        nll = 0.5 * (lv + (m - l).pow(2) / v).mean()
        L_tta_obs = mse + nll_weight * nll
    else:
        mse = zero; nll = zero; L_tta_obs = zero

    # ── TTA loss: ego_positive censored ──────────────────────────────────────
    if cens_mask.any():
        cm = tm[cens_mask]
        L_tta_cens = F.relu(tta_cap - cm).pow(2).mean()
    else:
        L_tta_cens = zero

    loss = L_hazard + tta_obs_weight * L_tta_obs + tta_cens_weight * L_tta_cens

    # ── metrics ──────────────────────────────────────────────────────────────
    n_obs   = int(obs_mask.sum())
    n_cens  = int(cens_mask.sum())
    n_pos   = n_obs + n_cens
    n_noneego = int((~is_ego_positive & (hazard_label == 0) & (hazard_weight < 0.5)).sum())

    hazard_pred_bin = (hp > 0.5).float()
    hazard_correct  = (hazard_pred_bin == hl).float().mean()

    metrics: Dict[str, float] = {
        "loss":             float(loss.detach()),
        "loss_hazard":      float(L_hazard.detach()),
        "loss_tta_obs":     float(L_tta_obs.detach()),
        "loss_tta_cens":    float(L_tta_cens.detach()),
        "hazard_acc":       float(hazard_correct),
        "n_obs":            n_obs,
        "n_cens":           n_cens,
        "n_pos":            n_pos,
        "n_noneego":        n_noneego,
    }
    if obs_mask.any():
        metrics["tta_mae"]  = float((tm[obs_mask] - tl[obs_mask]).abs().mean().detach())
        metrics["tta_rmse"] = float((tm[obs_mask] - tl[obs_mask]).pow(2).mean().sqrt().detach())
        metrics["mse_loss"] = float(mse.detach())
        metrics["nll_loss"] = float(nll.detach())
    else:
        metrics.update({"tta_mae": 0.0, "tta_rmse": 0.0, "mse_loss": 0.0, "nll_loss": 0.0})

    return loss, metrics


# ============================================================================
# Resume Helpers
# ============================================================================

def _is_sft_ckpt_dir(d: Path) -> bool:
    return (
        d.is_dir()
        and (d / "tta_head.pt").exists()
        and (d / "belief_aggregator.pt").exists()
        and (d / "config.json").exists()
        and (d / "vlm_lora" / "adapter_config.json").exists()
        and (d / "vlm_lora" / "adapter_model.safetensors").exists()
    )

def _parse_step(name: str) -> int:
    if name.startswith("step_"):
        try:
            return int(name.split("_", 1)[1])
        except Exception:
            return -1
    return -1

def find_auto_resume_checkpoint(output_dir: Path, experiment_name: str) -> Optional[Path]:
    candidates: List[Path] = []

    exp_dir = output_dir / experiment_name
    if exp_dir.exists():
        for child in exp_dir.iterdir():
            if _is_sft_ckpt_dir(child):
                candidates.append(child)

    if not candidates:
        for d1 in output_dir.iterdir():
            if not d1.is_dir():
                continue
            for d2 in d1.iterdir():
                if _is_sft_ckpt_dir(d2):
                    candidates.append(d2)

    if not candidates:
        return None

    step_cands = [(c, _parse_step(c.name)) for c in candidates]
    step_cands = [x for x in step_cands if x[1] >= 0]
    if step_cands:
        step_cands.sort(key=lambda x: x[1], reverse=True)
        return step_cands[0][0]

    epoch_cands = []
    for c in candidates:
        if c.name.startswith("epoch_"):
            try:
                epoch_cands.append((c, int(c.name.split("_", 1)[1])))
            except Exception:
                pass
    if epoch_cands:
        epoch_cands.sort(key=lambda x: x[1], reverse=True)
        return epoch_cands[0][0]

    for c in candidates:
        if c.name == "best":
            return c

    candidates.sort(key=lambda p: p.stat().st_mtime, reverse=True)
    return candidates[0]


def load_sft_heads(model: SFTModel, ckpt_dir: Path):
    b_path = ckpt_dir / "belief_aggregator.pt"
    h_path = ckpt_dir / "hazard_head.pt"
    t_path = ckpt_dir / "tta_head.pt"
    model.belief_aggregator.load_state_dict(torch.load(b_path, map_location=model.device), strict=True)
    if h_path.exists():
        model.hazard_head.load_state_dict(torch.load(h_path, map_location=model.device), strict=True)
        logger.info(f"  Loaded hazard_head from {h_path}")
    else:
        logger.warning("⚠️  hazard_head.pt not found in checkpoint; using fresh init.")
    model.tta_head.load_state_dict(torch.load(t_path, map_location=model.device), strict=True)
    logger.info(f"✅ Loaded heads from {ckpt_dir}")

    try:
        last = model.tta_head.net[-1]
        if hasattr(last, "bias") and last.bias is not None:
            logger.info(f"  TTAHead last-layer bias(after load) = {last.bias.detach().float().cpu().tolist()}")
    except Exception:
        pass





def compute_calibration_error(
    predictions: np.ndarray,
    uncertainties: np.ndarray,
    labels: np.ndarray,
    num_bins: int = 10
) -> Tuple[float, np.ndarray, np.ndarray]:
    """
    Compute Expected Calibration Error (ECE) for regression.

    Returns:
        ece: scalar
        observed_freq: per-bin observed frequencies
        expected_freq: per-bin expected frequencies
    """
    predictions = np.asarray(predictions)
    uncertainties = np.asarray(uncertainties)
    labels = np.asarray(labels)

    # normalized error = |pred-label| / std
    errors = np.abs(predictions - labels)
    normalized_errors = errors / np.maximum(uncertainties, 1e-6)

    if normalized_errors.size == 0:
        return 0.0, np.array([]), np.array([])

    # bin edges over normalized error
    max_ne = float(np.max(normalized_errors))
    if not np.isfinite(max_ne) or max_ne <= 0:
        return 0.0, np.array([]), np.array([])

    bin_edges = np.linspace(0.0, max_ne, num_bins + 1)

    observed_freq = []
    expected_freq = []

    # expected coverage for Gaussian within z std: erf(z/sqrt(2))
    # (note: this is a simple reference curve; you can replace later with your preferred ECE definition)
    sqrt2 = math.sqrt(2.0)

    for i in range(num_bins):
        lo, hi = bin_edges[i], bin_edges[i + 1]
        mask = (normalized_errors >= lo) & (normalized_errors < hi)
        if mask.sum() == 0:
            continue

        z = 0.5 * (lo + hi)
        expected = math.erf(z / sqrt2)  # in [0,1]
        observed = float(mask.mean())

        observed_freq.append(observed)
        expected_freq.append(expected)

    observed_freq = np.asarray(observed_freq, dtype=np.float32)
    expected_freq = np.asarray(expected_freq, dtype=np.float32)

    ece = float(np.abs(observed_freq - expected_freq).mean()) if observed_freq.size > 0 else 0.0
    return ece, observed_freq, expected_freq


# ============================================================================
# Trainer
# ============================================================================

class SFTTrainer:
    def __init__(
        self,
        model: SFTModel,
        train_dataset: SFTDataset,
        val_dataset: Optional[SFTDataset],
        num_epochs: int = 10,
        batch_size: int = 4,
        gradient_accumulation_steps: int = 4,
        learning_rate: float = 1e-4,
        tta_head_lr: float = 1e-3,
        vlm_lr_multiplier: float = 0.1,
        weight_decay: float = 0.01,
        max_grad_norm: float = 1.0,
        mse_weight: float = 1.0,
        nll_weight: float = 0.5,
        use_curriculum: bool = True,
        scheduler_type: str = "cosine",
        warmup_ratio: float = 0.1,
        output_dir: str = "./checkpoints/sft",
        experiment_name: str = "sft_default",
        logging_steps: int = 1250,
        eval_steps: int = 2500,
        save_steps: int = 5000,
        save_total_limit: int = 3,
        use_amp: bool = True,
        use_wandb: bool = True,
        wandb_project: str = "lkalert-sft",
        lora_update_patience: int = 30,
        disable_metadata_prompt: bool = False,   # P0.1: drop weather/road/time context
    ):
        self.model = model
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset

        self.num_epochs = num_epochs
        self.batch_size = batch_size
        self.gradient_accumulation_steps = gradient_accumulation_steps
        self.learning_rate = learning_rate
        self.tta_head_lr = tta_head_lr
        self.vlm_lr_multiplier = vlm_lr_multiplier
        self.weight_decay = weight_decay
        self.max_grad_norm = max_grad_norm
        self.mse_weight = mse_weight
        self.nll_weight = nll_weight
        self.use_curriculum = use_curriculum
        self.scheduler_type = scheduler_type
        self.warmup_ratio = warmup_ratio

        self.output_dir = Path(output_dir) / experiment_name
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.experiment_name = experiment_name
        self.logging_steps = logging_steps
        self.eval_steps = eval_steps
        self.save_steps = save_steps
        self.save_total_limit = save_total_limit

        # AMP
        self.use_amp = use_amp
        self.amp_dtype = torch.bfloat16 if self.model.dtype == torch.bfloat16 else torch.float16
        self.use_scaler = self.use_amp and (self.amp_dtype == torch.float16)
        self.scaler = GradScaler("cuda", enabled=self.use_scaler) if self.use_amp else None
        logger.info(f"AMP enabled={self.use_amp}, amp_dtype={self.amp_dtype}, scaler_enabled={self.use_scaler}")

        # wandb
        self.use_wandb = bool(use_wandb and HAS_WANDB)
        if self.use_wandb:
            wandb.init(
                project=wandb_project,
                name=experiment_name,
                config={
                    "num_epochs": num_epochs,
                    "batch_size": batch_size,
                    "grad_accum": gradient_accumulation_steps,
                    "learning_rate": learning_rate,
                    "tta_head_lr": tta_head_lr,
                    "vlm_lr_multiplier": vlm_lr_multiplier,
                    "use_curriculum": use_curriculum,
                },
            )

        # loaders/optim/sched
        self._create_dataloaders()
        self._create_optimizer()
        self._create_scheduler()

        # training state
        self.global_step = 0
        self.current_epoch = 0
        self.best_ckpt_score = float("-inf")   # higher is better (0.6*f1 - 0.4*mae/10)
        self.saved_checkpoints: List[Path] = []

        # LoRA checks
        self._lora_grad_verified = False
        self._lora_update_verified = False
        self._lora_update_zero_steps = 0
        self.lora_update_patience = int(lora_update_patience)

        # P0.1: if True, drop weather/road_type/time_of_day from the prompt.
        self.disable_metadata_prompt = bool(disable_metadata_prompt)

        logger.info("✅ SFTTrainer initialized")
        if self.disable_metadata_prompt:
            logger.info("  [P0.1] disable_metadata_prompt=True → metadata context stripped from prompt")
        logger.info(f"  Output dir: {self.output_dir}")
        logger.info(f"  Total steps: {self.total_steps}")
        logger.info(f"  Effective batch size: {batch_size * gradient_accumulation_steps}")

    def _create_dataloaders(self):
        self.train_loader = DataLoader(
            self.train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            collate_fn=sft_collate_fn,
            num_workers=4,
            pin_memory=True,
        )
        self.train_sampler = None

        self.val_loader = None
        if self.val_dataset is not None:
            self.val_loader = DataLoader(
                self.val_dataset,
                batch_size=self.batch_size * 2,
                shuffle=False,
                collate_fn=sft_collate_fn,
                num_workers=4,
                pin_memory=True,
            )

        steps_per_epoch = max(1, len(self.train_loader) // self.gradient_accumulation_steps)
        self.total_steps = steps_per_epoch * self.num_epochs

    def _create_optimizer(self):
        vlm_params = []
        for _, p in self.model.vlm.named_parameters():
            if p.requires_grad:
                vlm_params.append(p)

        head_params = (
            list(self.model.belief_aggregator.parameters())
            + list(self.model.hazard_head.parameters())
            + list(self.model.tta_head.parameters())
        )

        self.optimizer = AdamW(
            [
                {"params": vlm_params, "lr": self.learning_rate * self.vlm_lr_multiplier},
                {"params": head_params, "lr": self.tta_head_lr},
            ],
            weight_decay=self.weight_decay,
        )
        logger.info(f"  VLM params: {len(vlm_params)} (lr={self.learning_rate * self.vlm_lr_multiplier})")
        logger.info(f"  Head params: {len(head_params)} (lr={self.tta_head_lr})")

    def _create_scheduler(self):
        warmup_steps = int(self.total_steps * self.warmup_ratio)

        if self.scheduler_type == "cosine":
            warmup = LinearLR(self.optimizer, start_factor=0.1, end_factor=1.0, total_iters=max(1, warmup_steps))
            cosine = CosineAnnealingLR(self.optimizer, T_max=max(1, self.total_steps - warmup_steps), eta_min=1e-6)
            self.scheduler = SequentialLR(self.optimizer, schedulers=[warmup, cosine], milestones=[warmup_steps])
        elif self.scheduler_type == "linear":
            self.scheduler = LinearLR(self.optimizer, start_factor=1.0, end_factor=0.1, total_iters=max(1, self.total_steps))
        else:
            self.scheduler = None

    def _build_prompt(self, batch: Dict, idx: int) -> str:
        metadata = batch["metadata"][idx]
        window_type = batch["window_types"][idx]
        window_str = f"{2.0 if window_type == 'standard' else 3.0}s"

        # P0.1: prompt ablation — drop metadata context entirely.
        # Tests whether weather/road_type/time_of_day in the prompt contributes any
        # real signal vs. merely being boilerplate for the VLM.
        if getattr(self, "disable_metadata_prompt", False):
            return (
                f"Analyze this driving sequence ({window_str} window).\n"
                f"Estimate the time to potential collision. Output a single number in seconds."
            )

        context_parts = []
        if metadata.get("weather"):
            context_parts.append(f"Weather: {metadata['weather']}")
        if metadata.get("road_type"):
            context_parts.append(f"Road: {metadata['road_type']}")
        if metadata.get("time_of_day"):
            context_parts.append(f"Time: {metadata['time_of_day']}")
        context = ", ".join(context_parts) if context_parts else "Urban driving"

        return (
            f"Analyze this driving sequence ({window_str} window).\n"
            f"Context: {context}\n"
            f"Estimate the time to potential collision. Output a single number in seconds."
        )

    def _prepare_batch(self, batch: Dict) -> Dict[str, torch.Tensor]:
        system_prompt = "You are a driving safety AI analyzing dashcam footage for collision risk."

        texts = []
        images = batch["images"]  # List[List[PIL.Image]]: B x K
        proc = self.model.processor
        apply_chat = proc.apply_chat_template if hasattr(proc, "apply_chat_template") else proc.tokenizer.apply_chat_template

        for i in range(len(batch["video_ids"])):
            user_text = self._build_prompt(batch, i)
            frames = images[i]
            content = [{"type": "image"} for _ in range(len(frames))]
            content.append({"type": "text", "text": user_text})
            messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": content}]
            texts.append(apply_chat(messages, tokenize=False, add_generation_prompt=False))

        processed = proc(
            text=texts,
            images=images,
            return_tensors="pt",
            padding=True,
            truncation=True,
        )
        return processed

    # -------- LoRA checks --------

    def _verify_lora_grads_once(self):
        if self._lora_grad_verified:
            return
        lora = [(n, p) for n, p in self.model.vlm.named_parameters() if p.requires_grad and "lora_" in n.lower()]
        if not lora:
            logger.warning("⚠️ No trainable LoRA parameters found.")
            self._lora_grad_verified = True
            return
        non_none = 0
        non_zero = 0
        for _, p in lora:
            if p.grad is not None:
                non_none += 1
                if float(p.grad.detach().abs().sum().item()) > 0:
                    non_zero += 1
        logger.info(f"🔎 LoRA grad check: total={len(lora)}, grad_non_none={non_none}, grad_non_zero={non_zero}")
        if non_none == 0 or non_zero == 0:
            logger.warning("⚠️ LoRA grads are missing/zero at this moment (may be before first real update).")
        else:
            logger.info("✅ LoRA gradient flow verified.")
        self._lora_grad_verified = True

    def _pick_probe_lora_param(self) -> Optional[Tuple[str, torch.nn.Parameter]]:
        candidates = []
        for n, p in self.model.vlm.named_parameters():
            if not p.requires_grad:
                continue
            if "lora_" not in n.lower():
                continue
            if p.grad is None:
                continue
            if float(p.grad.detach().abs().sum().item()) == 0.0:
                continue
            candidates.append((n, p))
        if candidates:
            return random.choice(candidates)

        fallback = [(n, p) for n, p in self.model.vlm.named_parameters() if p.requires_grad and "lora_" in n.lower()]
        if not fallback:
            return None
        return random.choice(fallback)

    def _post_step_lora_update_check(self, probe_name: Optional[str], before_fp32: Optional[torch.Tensor]):
        if probe_name is None or before_fp32 is None:
            return

        probe_param = None
        for n, p in self.model.vlm.named_parameters():
            if n == probe_name:
                probe_param = p
                break
        if probe_param is None:
            return

        after_fp32 = probe_param.detach().float()
        delta = float((after_fp32 - before_fp32).abs().mean().item())

        if delta == 0.0:
            self._lora_update_zero_steps += 1
            lr0 = self.optimizer.param_groups[0]["lr"]
            logger.warning(
                f"⚠️ LoRA update probe delta==0 (name='{probe_name}'), "
                f"consecutive_zero_steps={self._lora_update_zero_steps}, lr={lr0:.2e}. "
                f"Will only raise after {self.lora_update_patience} consecutive steps."
            )
            if self._lora_update_zero_steps >= self.lora_update_patience:
                raise RuntimeError(
                    "LoRA probe parameter did not change for many optimizer steps. "
                    "Likely lr too small for bf16 rounding, or LoRA params not in optimizer, or training graph bypassing LoRA."
                )
        else:
            if not self._lora_update_verified:
                logger.info(f"✅ LoRA update verified: probe='{probe_name}', mean_abs_delta={delta:.6e}")
            self._lora_update_verified = True
            self._lora_update_zero_steps = 0

    # -------- train/eval --------

    def _batch_to_device(self, batch: Dict, keys) -> Dict:
        return {k: batch[k].to(self.model.device) for k in keys if k in batch}

    def train_step(self, batch: Dict) -> Dict[str, float]:
        self.model.train()
        inputs = self._prepare_batch(batch)
        t = self._batch_to_device(batch, [
            "tta_labels", "hazard_labels", "hazard_weights",
            "is_ego_positive", "is_censored",
        ])

        with autocast(device_type="cuda", dtype=self.amp_dtype, enabled=self.use_amp):
            out = self.model(inputs)
            loss, metrics = compute_sft_loss(
                hazard_logit    = out["hazard_logit"],
                tta_mean        = out["tta_mean"],
                tta_logvar      = out["tta_logvar"],
                hazard_label    = t["hazard_labels"],
                hazard_weight   = t["hazard_weights"],
                is_ego_positive = t["is_ego_positive"],
                is_censored     = t["is_censored"],
                tta_label       = t["tta_labels"],
                nll_weight      = self.nll_weight,
            )
            loss = loss / self.gradient_accumulation_steps

        if self.use_scaler:
            self.scaler.scale(loss).backward()
        else:
            loss.backward()

        if not self._lora_grad_verified:
            self._verify_lora_grads_once()

        return metrics

    def _optimizer_step(self):
        probe = self._pick_probe_lora_param()
        probe_name = probe[0] if probe else None
        before_fp32 = probe[1].detach().float().clone() if probe else None

        if self.use_scaler:
            self.scaler.unscale_(self.optimizer)

        torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)

        if self.use_scaler:
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            self.optimizer.step()

        self.optimizer.zero_grad(set_to_none=True)
        if self.scheduler is not None:
            self.scheduler.step()

        self._post_step_lora_update_check(probe_name, before_fp32)

        self.global_step += 1

    @torch.no_grad()
    def evaluate(self) -> Dict[str, float]:
        if self.val_loader is None:
            return {}
        self.model.eval()

        total_loss = 0.0
        n = 0
        preds, labels_all, stds = [], [], []

        all_hazard_prob:    List[np.ndarray] = []
        all_hazard_label:   List[np.ndarray] = []
        all_is_noneego:     List[np.ndarray] = []
        all_is_ego_pos:     List[np.ndarray] = []

        for batch in tqdm(self.val_loader, desc="Evaluating", leave=False, ncols=60):
            inputs = self._prepare_batch(batch)
            t = self._batch_to_device(batch, [
                "tta_labels", "hazard_labels", "hazard_weights",
                "is_ego_positive", "is_censored",
            ])
            is_noneego_b = batch.get("is_non_ego", torch.zeros(len(batch["video_ids"]), dtype=torch.bool))

            with autocast(device_type="cuda", dtype=self.amp_dtype, enabled=self.use_amp):
                out = self.model(inputs)
                loss, _ = compute_sft_loss(
                    hazard_logit    = out["hazard_logit"],
                    tta_mean        = out["tta_mean"],
                    tta_logvar      = out["tta_logvar"],
                    hazard_label    = t["hazard_labels"],
                    hazard_weight   = t["hazard_weights"],
                    is_ego_positive = t["is_ego_positive"],
                    is_censored     = t["is_censored"],
                    tta_label       = t["tta_labels"],
                    nll_weight      = self.nll_weight,
                )

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

            tta_mean  = out["tta_mean"].detach().float().cpu().numpy()
            tta_label_np = t["tta_labels"].detach().float().cpu().numpy()
            tta_std   = torch.exp(0.5 * out["tta_logvar"].detach().float()).cpu().numpy()

            preds.append(tta_mean)
            labels_all.append(tta_label_np)
            stds.append(tta_std)

            all_hazard_prob.append(out["hazard_prob"].detach().float().cpu().numpy())
            all_hazard_label.append(t["hazard_labels"].detach().float().cpu().numpy())
            all_is_noneego.append(is_noneego_b.cpu().numpy())
            all_is_ego_pos.append(t["is_ego_positive"].cpu().numpy())

        preds      = np.concatenate(preds)       if preds       else np.zeros(0, np.float32)
        labels_all = np.concatenate(labels_all)  if labels_all  else np.zeros(0, np.float32)
        stds       = np.concatenate(stds)        if stds        else np.zeros(0, np.float32)
        hp_all     = np.concatenate(all_hazard_prob)  if all_hazard_prob  else np.zeros(0, np.float32)
        hl_all     = np.concatenate(all_hazard_label) if all_hazard_label else np.zeros(0, np.float32)
        ne_all     = np.concatenate(all_is_noneego)   if all_is_noneego   else np.zeros(0, bool)
        ep_all     = np.concatenate(all_is_ego_pos)   if all_is_ego_pos   else np.zeros(0, bool)

        if preds.size == 0:
            self.model.train()
            return {"loss": float("inf"), "hazard_f1": 0.0, "pos_tta_mae": float("inf"),
                    "ckpt_score": float("-inf")}

        # ── hazard metrics ───────────────────────────────────────────────────
        hp_bin = (hp_all > 0.5).astype(np.float32)
        tp = float(((hp_bin == 1) & (hl_all == 1)).sum())
        fp = float(((hp_bin == 1) & (hl_all == 0)).sum())
        fn = float(((hp_bin == 0) & (hl_all == 1)).sum())
        prec   = tp / max(1, tp + fp)
        recall = tp / max(1, tp + fn)
        f1     = 2 * prec * recall / max(1e-9, prec + recall)

        ne_mask      = ne_all.astype(bool)
        safe_neg_mask = (~ep_all) & (~ne_mask)
        ne_far  = float((hp_bin[ne_mask] == 1).mean())     if ne_mask.any()       else 0.0
        sneg_fa = float((hp_bin[safe_neg_mask] == 1).mean()) if safe_neg_mask.any() else 0.0

        # ── TTA metrics (positive-observed only) ─────────────────────────────
        obs_mask = ep_all & (labels_all < 9.9)
        if obs_mask.any():
            pos_preds  = preds[obs_mask];  pos_labels = labels_all[obs_mask]
            pos_mae    = float(np.abs(pos_preds - pos_labels).mean())
            pos_rmse   = float(np.sqrt(((pos_preds - pos_labels)**2).mean()))
            low_mask   = pos_labels <= 3.0
            low_mae    = float(np.abs(pos_preds[low_mask] - pos_labels[low_mask]).mean()) if low_mask.any() else 0.0
            denom = float(((pos_labels - pos_labels.mean())**2).sum()) + 1e-12
            pos_r2 = float(1.0 - ((pos_preds - pos_labels)**2).sum() / denom)
        else:
            pos_mae = pos_rmse = low_mae = 10.0; pos_r2 = 0.0

        # ── checkpoint selection score ────────────────────────────────────────
        # Higher is better: maximize hazard_f1, minimize normalized pos_tta_mae
        ckpt_score = 0.6 * f1 - 0.4 * (pos_mae / 10.0)

        metrics = {
            "loss":                  total_loss / max(1, n),
            "hazard_f1":             f1,
            "hazard_precision":      prec,
            "hazard_recall":         recall,
            "pos_tta_mae":           pos_mae,
            "pos_tta_rmse":          pos_rmse,
            "pos_tta_r2":            pos_r2,
            "low_tta_mae":           low_mae,
            "non_ego_false_alert":   ne_far,
            "safe_neg_false_alert":  sneg_fa,
            "uncertainty_mean":      float(stds.mean()),
            "ckpt_score":            ckpt_score,
        }

        logger.info(
            f"Val: loss={metrics['loss']:.4f}  hazard_f1={f1:.3f}  "
            f"pos_tta_mae={pos_mae:.3f}  ckpt_score={ckpt_score:.4f}  "
            f"non_ego_fa={ne_far:.3f}  safe_neg_fa={sneg_fa:.3f}"
        )
        self.model.train()
        return metrics

    def save_checkpoint(self, name: str):
        ckpt_dir = self.output_dir / name
        self.model.save_checkpoint(str(ckpt_dir), epoch=self.current_epoch, step=self.global_step)

        torch.save(
            {
                "optimizer": self.optimizer.state_dict(),
                "scheduler": self.scheduler.state_dict() if self.scheduler else None,
                "scaler": self.scaler.state_dict() if self.scaler else None,
                "epoch": self.current_epoch,
                "global_step": self.global_step,
                "best_ckpt_score": self.best_ckpt_score,
            },
            ckpt_dir / "training_state.pt",
        )

        self.saved_checkpoints.append(ckpt_dir)
        if len(self.saved_checkpoints) > self.save_total_limit + 1:
            oldest = self.saved_checkpoints.pop(0)
            if oldest.name != "best" and oldest.exists():
                import shutil
                shutil.rmtree(oldest, ignore_errors=True)

    def load_training_state(self, ckpt_dir: Path, reset_best_val_loss: bool = False):
        """
        Loads optimizer/scheduler/scaler + epoch/global_step.
        If reset_best_val_loss=True: best_ckpt_score is forcibly reset to -inf
        so your NEW val split can define a NEW best from scratch.
        """
        ts = ckpt_dir / "training_state.pt"
        if not ts.exists():
            logger.warning(f"⚠️ No training_state.pt in {ckpt_dir}, resume weights only.")
            if reset_best_val_loss:
                self.best_ckpt_score = float("-inf")
                logger.info("✅ best_ckpt_score reset to -inf (weights-only path).")
            return

        obj = torch.load(ts, map_location="cpu")
        try:
            self.optimizer.load_state_dict(obj["optimizer"])
        except Exception as e:
            logger.warning(f"⚠️ Failed to load optimizer state: {e}")
        if self.scheduler is not None and obj.get("scheduler") is not None:
            try:
                self.scheduler.load_state_dict(obj["scheduler"])
            except Exception as e:
                logger.warning(f"⚠️ Failed to load scheduler state: {e}")
        if self.scaler is not None and obj.get("scaler") is not None:
            try:
                self.scaler.load_state_dict(obj["scaler"])
            except Exception as e:
                logger.warning(f"⚠️ Failed to load scaler state: {e}")

        self.current_epoch = int(obj.get("epoch", 0))
        self.global_step = int(obj.get("global_step", 0))

        if reset_best_val_loss:
            self.best_ckpt_score = float("-inf")
            logger.info(
                f"✅ Loaded training_state from {ckpt_dir}: epoch={self.current_epoch}, global_step={self.global_step}, "
                f"best_ckpt_score RESET to -inf for NEW val."
            )
        else:
            self.best_ckpt_score = float(obj.get("best_ckpt_score", obj.get("best_val_loss", float("-inf"))))
            logger.info(
                f"✅ Loaded training_state from {ckpt_dir}: epoch={self.current_epoch}, global_step={self.global_step}, "
                f"best_ckpt_score={self.best_ckpt_score:.4f}"
            )

    def maybe_eval_and_set_new_best(self, force_save_best: bool = True):
        """
        Evaluate once immediately (useful after resume + reset_best_val_loss).
        If best_ckpt_score is -inf, this will always become the new best.
        """
        if self.val_loader is None:
            return

        val = self.evaluate()
        if self.use_wandb and val:
            wandb.log({"val/" + k: v for k, v in val.items()}, step=self.global_step)

        if not val:
            return

        score = val.get("ckpt_score", float("-inf"))
        improved = score > self.best_ckpt_score
        if improved:
            self.best_ckpt_score = score
            if force_save_best:
                self.save_checkpoint("best")

        logger.info(
            f"[InitEval] ckpt_score={score:.4f}, "
            f"best_ckpt_score={self.best_ckpt_score:.4f}, improved={improved}"
        )

    def train(self):
        logger.info("=" * 60)
        logger.info(f"Starting SFT training: {self.experiment_name}")
        logger.info("=" * 60)

        start = time.time()

        for epoch in range(self.current_epoch, self.num_epochs):
            self.current_epoch = epoch

            progress = epoch / max(1, self.num_epochs)
            _ = progress  # curriculum not used in new dataset

            pbar = tqdm(self.train_loader, desc=f"Epoch {epoch+1}/{self.num_epochs}", ncols=60)
            metrics_hist = defaultdict(list)
            accum = 0

            for batch in pbar:
                m = self.train_step(batch)
                for k, v in m.items():
                    metrics_hist[k].append(v)
                accum += 1

                if accum >= self.gradient_accumulation_steps:
                    self._optimizer_step()
                    accum = 0

                    if self.global_step % self.logging_steps == 0:
                        avg = {k: float(np.mean(v[-self.logging_steps:])) for k, v in metrics_hist.items()}
                        lr = self.optimizer.param_groups[0]["lr"]
                        pbar.set_postfix({"loss": f"{avg['loss']:.4f}", "mae": f"{avg['tta_mae']:.3f}", "lr": f"{lr:.2e}"})
                        if self.use_wandb:
                            wandb.log({"train/" + k: v for k, v in avg.items()} | {"train/lr": lr}, step=self.global_step)

                    if self.val_loader and (self.global_step % self.eval_steps == 0):
                        val = self.evaluate()
                        if self.use_wandb and val:
                            wandb.log({"val/" + k: v for k, v in val.items()}, step=self.global_step)
                        if val:
                            score = val.get("ckpt_score", float("-inf"))
                            if score > self.best_ckpt_score:
                                self.best_ckpt_score = score
                                self.save_checkpoint("best")

                    if self.global_step % self.save_steps == 0:
                        self.save_checkpoint(f"step_{self.global_step}")

            if self.val_loader:
                val = self.evaluate()
                if val:
                    score = val.get("ckpt_score", float("-inf"))
                    if score > self.best_ckpt_score:
                        self.best_ckpt_score = score
                        self.save_checkpoint("best")
            self.save_checkpoint(f"epoch_{epoch+1}")

        logger.info("=" * 60)
        logger.info(f"Training completed in {(time.time()-start)/3600:.2f} hours")
        logger.info(f"Best ckpt_score: {self.best_ckpt_score:.4f}")
        logger.info(f"Checkpoints saved to: {self.output_dir}")
        logger.info("=" * 60)

        if self.use_wandb:
            wandb.finish()


# ============================================================================
# Main
# ============================================================================

def main():
    parser = argparse.ArgumentParser("SFT Training for TTA Regression")

    # data — manifest-based
    parser.add_argument(
        "--manifest_dir", type=str,
        default="PROJECT_ROOT/data/sft_manifests",
        help="Directory containing split manifest JSONs from make_split_manifest.py",
    )
    # legacy aliases (ignored if manifest_dir is provided via manifests)
    parser.add_argument("--nexar_root", type=str, default=None, help="(unused; kept for back-compat)")
    parser.add_argument("--dada_root", type=str, default=None, help="(unused; kept for back-compat)")

    # model
    parser.add_argument("--model_name", type=str, default="Qwen/Qwen2.5-VL-3B-Instruct")
    parser.add_argument("--pretrained_lora", type=str, default=None)
    parser.add_argument(
        "--attn_implementation", type=str, default="flash_attention_2",
        choices=["flash_attention_2", "sdpa", "eager"],
        help="VLM attention backend. sdpa is safe fallback for Blackwell/new backbones.",
    )
    parser.add_argument(
        "--belief_strategy", type=str, default="mean_pool",
        choices=["mean_pool", "last_token", "attention_pool", "dual_pool"],
        help="dual_pool = [mean(image_tokens) || mean(text_tokens)]  (P0.2 L1)",
    )

    # P0.1 — prompt ablation (drop weather/road_type/time_of_day)
    parser.add_argument(
        "--disable_metadata_prompt", action="store_true", default=False,
        help="P0.1: remove weather/road_type/time_of_day from the SFT prompt.",
    )

    # P0.3 — PEFT upgrade flags
    parser.add_argument("--use_dora",   action="store_true", default=False,
                        help="P0.3: enable DoRA (Weight-Decomposed Low-Rank Adaptation).")
    parser.add_argument("--use_rslora", action="store_true", default=False,
                        help="P0.3: enable rsLoRA (rank-stabilised scaling alpha/sqrt(r)).")
    parser.add_argument(
        "--lora_init", type=str, default="default",
        choices=["default", "gaussian", "pissa", "pissa_niter_16", "olora"],
        help="P0.3: initialisation scheme for fresh LoRA (ignored when resuming).",
    )

    # resume
    parser.add_argument("--resume_from", type=str, default=None,
                        help="Path to an SFT checkpoint dir that contains tta_head.pt, belief_aggregator.pt, and vlm_lora/")
    parser.add_argument("--resume_weights_only", action="store_true",
                        help="If set, do not load optimizer/scheduler/scaler states (start new training state).")
    parser.add_argument("--auto_resume", action="store_true", default=True,
                        help="Auto search a previous SFT checkpoint under output_dir (default: True).")
    parser.add_argument("--no_auto_resume", action="store_false", dest="auto_resume")

    # NEW: reset best + optional eval at start
    parser.add_argument("--reset_best_val_loss", action="store_true", default=True,
                        help="Reset best_val_loss to +inf when resuming so NEW val can redefine best. (default: True)")
    parser.add_argument("--no_reset_best_val_loss", action="store_false", dest="reset_best_val_loss")
    parser.add_argument("--eval_on_start", action="store_true", default=True,
                        help="Run a val evaluation immediately after resume (useful with reset_best_val_loss). (default: True)")
    parser.add_argument("--no_eval_on_start", action="store_false", dest="eval_on_start")

    # training
    parser.add_argument("--num_epochs", type=int, default=10)
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=4)
    parser.add_argument("--learning_rate", type=float, default=1e-4)
    parser.add_argument("--tta_head_lr", type=float, default=1e-3)
    parser.add_argument("--vlm_lr_multiplier", type=float, default=0.1)
    parser.add_argument("--weight_decay", type=float, default=0.01)
    parser.add_argument("--max_grad_norm", type=float, default=1.0)
    parser.add_argument("--mse_weight", type=float, default=1.0)
    parser.add_argument("--nll_weight", type=float, default=0.5)
    parser.add_argument("--max_pixels", type=int, default=None,
                        help="Max pixels per frame for vision encoder. Default: 768*28*28=602112. "
                             "Lower (e.g. 512*28*28=401408) reduces VRAM → allows larger batch.")

    parser.add_argument("--use_curriculum", action="store_true", default=True)
    parser.add_argument("--no_curriculum", action="store_false", dest="use_curriculum")

    # output/log
    parser.add_argument("--output_dir", type=str, required=True)
    parser.add_argument("--experiment_name", type=str, required=True)
    parser.add_argument("--use_wandb", action="store_true", default=True)
    parser.add_argument("--no_wandb", action="store_false", dest="use_wandb")

    # debug
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--debug_samples", type=int, default=100)

    args = parser.parse_args()

    # datasets — manifest-based
    logger.info("📊 Loading datasets from manifests...")
    manifest_dir = Path(args.manifest_dir)

    train_manifests = [
        manifest_dir / "nexar_train.json",
        manifest_dir / "dada_pos_train.json",
        manifest_dir / "dada_noneego_train.json",
        manifest_dir / "dada_neg_train.json",
    ]
    val_manifests = [
        manifest_dir / "nexar_val.json",
        manifest_dir / "dada_pos_val.json",
        manifest_dir / "dada_noneego_val.json",
    ]

    # Filter to existing manifests (graceful in case some sources are absent)
    train_manifests = [m for m in train_manifests if m.exists()]
    val_manifests   = [m for m in val_manifests   if m.exists()]

    if not train_manifests:
        raise RuntimeError(f"No train manifests found in {manifest_dir}. Run make_split_manifest.py first.")

    logger.info(f"  Train manifests: {[m.name for m in train_manifests]}")
    logger.info(f"  Val   manifests: {[m.name for m in val_manifests]}")

    train_dataset = SFTDataset(
        manifests=train_manifests,
        split="train",
        debug=args.debug,
        debug_samples=args.debug_samples,
    )
    val_dataset = SFTDataset(
        manifests=val_manifests,
        split="val",
        debug=args.debug,
        debug_samples=max(1, args.debug_samples // 2),
    ) if val_manifests else None

    # Decide resume checkpoint
    output_root = Path(args.output_dir)
    resume_dir: Optional[Path] = None
    if args.resume_from:
        resume_dir = Path(args.resume_from)
        if not _is_sft_ckpt_dir(resume_dir):
            raise RuntimeError(f"--resume_from is not a valid SFT checkpoint dir: {resume_dir}")
    elif args.auto_resume:
        resume_dir = find_auto_resume_checkpoint(output_root, args.experiment_name)
        if resume_dir is not None:
            logger.info(f"🔁 Auto-resume selected checkpoint: {resume_dir}")

    # If resume: load LoRA from ckpt/vlm_lora
    lora_path_for_init = args.pretrained_lora
    if resume_dir is not None:
        lora_path_for_init = str(resume_dir / "vlm_lora")

    # Create model
    logger.info("📦 Creating model...")
    model = SFTModel(
        model_name=args.model_name,
        pretrained_lora_path=lora_path_for_init,
        belief_strategy=args.belief_strategy,
        use_lora=True,
        use_bf16=True,
        device="auto",
        max_pixels=args.max_pixels,
        use_dora=args.use_dora,
        use_rslora=args.use_rslora,
        lora_init=args.lora_init,
        attn_implementation=args.attn_implementation,
    )

    # If resume: load heads
    if resume_dir is not None:
        load_sft_heads(model, resume_dir)

    # Trainer
    trainer = SFTTrainer(
        model=model,
        train_dataset=train_dataset,
        val_dataset=val_dataset,
        num_epochs=args.num_epochs,
        batch_size=args.batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        tta_head_lr=args.tta_head_lr,
        vlm_lr_multiplier=args.vlm_lr_multiplier,
        weight_decay=args.weight_decay,
        max_grad_norm=args.max_grad_norm,
        mse_weight=args.mse_weight,
        nll_weight=args.nll_weight,
        use_curriculum=args.use_curriculum,
        output_dir=args.output_dir,
        experiment_name=args.experiment_name,
        use_wandb=args.use_wandb and HAS_WANDB,
        disable_metadata_prompt=args.disable_metadata_prompt,
    )

    # Load training state if requested (but allow best reset)
    if resume_dir is not None and (not args.resume_weights_only):
        trainer.load_training_state(resume_dir, reset_best_val_loss=args.reset_best_val_loss)
    else:
        # weights-only path: still respect reset_best_val_loss
        if resume_dir is not None and args.reset_best_val_loss:
            trainer.best_ckpt_score = float("-inf")
            logger.info("✅ best_ckpt_score reset to -inf (resume_weights_only path).")

    # IMPORTANT: Evaluate immediately on the NEW val, so best is re-defined from scratch.
    if resume_dir is not None and args.eval_on_start and trainer.val_loader is not None:
        # If best is inf, this will always save "best" corresponding to the NEW val baseline.
        trainer.maybe_eval_and_set_new_best(force_save_best=True)

    trainer.train()


if __name__ == "__main__":
    main()