"""
VLM预训练配置
支持多个模型和多任务学习
"""

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

@dataclass
class ModelConfig:
    """模型配置"""
    model_name: str = "Qwen2.5-VL-3B-Instruct"
    model_path: str = "PROJECT_ROOT/models/Qwen2.5-VL-3B-Instruct"
    model_type: str = "qwen2.5-vl"  # qwen2.5-vl, llava-onevision, minicpm-v, etc.
    
    # LoRA配置
    use_lora: bool = True
    lora_r: int = 32
    lora_alpha: int = 32
    lora_dropout: float = 0.1
    lora_target_modules: List[str] = field(default_factory=lambda: [
        "q_proj", "v_proj", "k_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ])
    
    # 量化
    load_in_4bit: bool = False
    load_in_8bit: bool = False


@dataclass
class DataConfig:
    """数据配置"""
    data_file: str = "PROJECT_ROOT/data/dataset/pretrain/train/pretrain_data.pkl"
    image_size: int = 224
    max_sequence_length: int = 30  # 任务3最大序列长度
    
    # 任务权重
    task1_weight: float = 1.0  # 环境描述
    task2_weight: float = 1.0  # 事故检测
    task3_weight: float = 2.0  # 序列预测（更重要）


@dataclass
class TrainingConfig:
    """训练配置"""
    output_dir: str = "PROJECT_ROOT/checkpoints/pretrain"
    
    # 训练参数
    num_epochs: int = 5
    batch_size: int = 4
    gradient_accumulation_steps: int = 4
    learning_rate: float = 2e-5
    weight_decay: float = 0.01
    warmup_ratio: float = 0.1
    max_grad_norm: float = 1.0
    
    # 优化器
    optimizer_type: str = "adamw"
    lr_scheduler_type: str = "cosine"
    
    # 日志和保存
    logging_steps: int = 10
    save_steps: int = 500
    save_total_limit: int = 3
    eval_steps: int = 500
    
    # 设备
    device: str = "cuda"
    fp16: bool = True
    bf16: bool = False
    
    # 随机种子
    seed: int = 42
    
    # wandb
    use_wandb: bool = False
    wandb_project: str = "lkalert-pretrain"
    wandb_run_name: Optional[str] = None


@dataclass
class PretrainConfig:
    """完整配置"""
    model: ModelConfig = field(default_factory=ModelConfig)
    data: DataConfig = field(default_factory=DataConfig)
    training: TrainingConfig = field(default_factory=TrainingConfig)
    
    def __post_init__(self):
        # 根据模型名称设置输出目录
        self.training.output_dir = os.path.join(
            self.training.output_dir,
            self.model.model_name
        )
        os.makedirs(self.training.output_dir, exist_ok=True)


# 预定义配置
QWEN25_VL_3B_CONFIG = PretrainConfig(
    model=ModelConfig(
        model_name="Qwen2.5-VL-3B-Instruct",
        model_path="PROJECT_ROOT/models/Qwen2.5-VL-3B-Instruct",
        model_type="qwen2.5-vl",
        lora_r=32,
        lora_alpha=32
    ),
    training=TrainingConfig(
        # batch_size=8,
        # gradient_accumulation_steps=2,
        batch_size=1,
        gradient_accumulation_steps=8,
        num_epochs=5
    )
)

QWEN25_VL_7B_CONFIG = PretrainConfig(
    model=ModelConfig(
        model_name="Qwen2.5-VL-7B-Instruct",
        model_path="PROJECT_ROOT/models/Qwen2.5-VL-7B-Instruct",
        model_type="qwen2.5-vl",
        lora_r=32,
        lora_alpha=32,
        load_in_8bit=True  # 7B模型使用8bit量化
    ),
    training=TrainingConfig(
        batch_size=4,
        gradient_accumulation_steps=4,
        num_epochs=5
    )
)