"""
配置管理
"""

from dataclasses import dataclass, field
from typing import Optional

@dataclass
class ModelConfig:
    """模型配置"""
    # VLM backbone
    model_name: str = "./models/Qwen2.5-VL-3B-Instruct"
    
    # 组件配置
    # 注意：不同模型的hidden_dim不同
    # Qwen2.5-VL-3B: 2048
    # Qwen2.5-VL-7B: 3584  
    # Qwen3-VL-4B: 2560
    tta_intermediate_dim: int = 512
    
    # belief聚合方式
    belief_aggregation: str = "mean_pool"  # "mean_pool" | "belief_token" | "attention_pool"
    
    # LoRA配置（可选）
    use_lora: bool = False
    lora_r: int = 32
    lora_alpha: int = 32
    lora_dropout: float = 0.1
    lora_target_modules: list = field(default_factory=lambda: [
        'q_proj', 'v_proj', 'k_proj', 'o_proj',
        'gate_proj', 'up_proj', 'down_proj'
    ])

@dataclass
class TrainingConfig:
    """训练配置"""
    # 基础设置
    output_dir: str = "./checkpoints/sft"
    num_epochs: int = 10
    batch_size: int = 4
    gradient_accumulation_steps: int = 4
    learning_rate: float = 2e-5
    weight_decay: float = 0.01
    warmup_steps: int = 1000
    max_grad_norm: float = 1.0
    
    # 损失权重
    lambda_nll: float = 0.5
    
    # Curriculum
    curriculum_warmup_ratio: float = 0.3
    curriculum_transition_ratio: float = 0.4
    
    # 保存和日志
    save_steps: int = 500
    logging_steps: int = 100
    eval_steps: int = 500
    save_total_limit: int = 3
    
    # 早停
    early_stopping_patience: int = 3
    early_stopping_metric: str = "val_mse"
    
    # 混合精度
    fp16: bool = False
    bf16: bool = True  # Qwen2.5-VL推荐使用bf16
    
    # DeepSpeed（可选）
    use_deepspeed: bool = False
    deepspeed_config: Optional[str] = None

@dataclass
class DataConfig:
    """数据配置"""
    # 数据路径
    train_data_path: str = "./data/processed/train/"
    val_data_path: str = "./data/processed/val/"
    
    # 视频参数
    video_window: float = 2.0  # 秒
    video_fps: int = 10
    video_height: int = 224
    video_width: int = 448
    max_frames: int = 20  # video_window * video_fps
    
    # 数据加载
    num_workers: int = 4
    pin_memory: bool = True
    prefetch_factor: int = 2
    
    # 数据增强
    use_augmentation: bool = True
    time_jitter: float = 0.2  # 时间抖动范围（秒）
    color_jitter: bool = True