#!/usr/bin/env python3
"""
Pretrain V2 — Path C configuration
====================================
Stage-A : BDD100K  → calibrated risk vocabulary (risk 1-2/5, no binary crash bias)
Stage-B : DADA-2000 + NEXAR  → TTA-labeled 2s windows (matches SFT inference)
"""

from dataclasses import dataclass, field
from typing import List, Optional

# ── Absolute paths ────────────────────────────────────────────────────────────
MODEL_PATH         = "PROJECT_ROOT/models/Qwen2.5-VL-3B-Instruct"
BDD100K_IMAGES_DIR = "PROJECT_ROOT/BDD-100K/bdd100k/images/100k"
BDD100K_LABELS_DIR = "PROJECT_ROOT/BDD-100K/bdd100k/labels/100k"
NEXAR_DATASET_DIR  = "PROJECT_ROOT/NEXAR_COLLISION/dataset"
DADA_DATASET_DIR   = "PROJECT_ROOT/DADA-2000"
DATA_OUTPUT_DIR    = "PROJECT_ROOT/data/pretrain_v2"
CKPT_BASE_DIR      = "PROJECT_ROOT/checkpoints/pretrain_v2"
STAGE_A_CKPT_DIR   = f"{CKPT_BASE_DIR}/stage_a"
STAGE_B_CKPT_DIR   = f"{CKPT_BASE_DIR}/stage_b"

# ── Generated data files ──────────────────────────────────────────────────────
STAGE_A_TRAIN_JSON = f"{DATA_OUTPUT_DIR}/stage_a_train.json"
STAGE_A_VAL_JSON   = f"{DATA_OUTPUT_DIR}/stage_a_val.json"
STAGE_B_TRAIN_JSON = f"{DATA_OUTPUT_DIR}/stage_b_train.json"
STAGE_B_VAL_JSON   = f"{DATA_OUTPUT_DIR}/stage_b_val.json"

# ── TTA clipping (matches SFT) ────────────────────────────────────────────────
TTA_MIN = 0.1
TTA_MAX = 10.0


def tta_to_risk(tta_s: float) -> int:
    """Map TTA in seconds to risk level 1-5."""
    if tta_s < 1.0:
        return 5
    if tta_s < 2.0:
        return 4
    if tta_s < 4.0:
        return 3
    if tta_s < 6.0:
        return 2
    return 1


# ── Data preparation config ───────────────────────────────────────────────────
@dataclass
class DataPrepConfig:
    # Stage-A BDD100K
    stage_a_max_per_task: int = 25_000   # 25k × 3 tasks = 75k training samples
    stage_a_val_ratio: float = 0.05

    # Stage-B TTA windows
    tta_deltas: List[float] = field(
        default_factory=lambda: [0.5, 1.0, 1.5, 2.0, 3.0, 4.5, 6.0]
    )
    window_duration_s: float = 2.0       # 2s window → 40 frames at 20fps
    n_frames_per_window: int = 8         # evenly sampled from 2s window
    dada_conservative_shift_s: float = 1.0  # paper §4.4: DADA annotations conservative
    stage_b_val_ratio: float = 0.10
    seed: int = 42


# ── LoRA config ───────────────────────────────────────────────────────────────
@dataclass
class LoraConfig:
    r: int = 32
    alpha: int = 32
    dropout: float = 0.05
    target_modules: List[str] = field(default_factory=lambda: [
        "q_proj", "v_proj", "k_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ])


# ── Training config (shared, override per stage) ──────────────────────────────
@dataclass
class TrainConfig:
    # Model
    model_path: str = MODEL_PATH
    lora: LoraConfig = field(default_factory=LoraConfig)
    # Processor pixel budgets.  Stage-A (single image): keep default.
    # Stage-B (8 frames): reduce to fit within GPU memory.
    #   128×28×28=100352 px/frame → ~120 tokens/frame × 8 frames ≈ 960 image tokens
    max_pixels_single:   int = 768 * 28 * 28   # Stage-A: one image, can afford high res
    max_pixels_sequence: int = 128 * 28 * 28   # Stage-B: 8 frames, must be small

    # Loop
    num_epochs: int = 1
    batch_size: int = 1
    gradient_accumulation_steps: int = 8
    learning_rate: float = 2e-5
    weight_decay: float = 0.01
    warmup_ratio: float = 0.05
    max_grad_norm: float = 1.0

    # Logging / saving
    logging_steps: int = 20
    eval_steps: int = 500
    save_steps: int = 500
    save_total_limit: int = 2

    # AMP
    bf16: bool = True

    # Wandb
    use_wandb: bool = True
    wandb_project: str = "lkalert-pretrain-v2"
    wandb_run_name: Optional[str] = None

    # Paths (set by train_stage_*.py)
    output_dir: str = CKPT_BASE_DIR
    pretrained_lora_path: Optional[str] = None   # Stage-B: path to Stage-A best_model