#!/usr/bin/env python3
"""
VLM预训练主脚本
支持多模型和多任务学习
"""

import os
import sys
import torch
import random
import numpy as np
import argparse
from torch.utils.data import DataLoader

# 添加路径
sys.path.insert(0, 'PROJECT_ROOT/data/dataset/pretrain')
from pretrain_dataset import PretrainDataset, collate_fn

from config import QWEN25_VL_3B_CONFIG, QWEN25_VL_7B_CONFIG
from trainer import MultiTaskTrainer


def set_seed(seed: int):
    """设置随机种子"""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def create_dataloaders(config):
    """创建数据加载器"""
    print("=" * 60)
    print("准备数据...")
    
    train_dataset = PretrainDataset(
        data_file=config.data.data_file,
        split="train",
        task="all"
    )
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=config.training.batch_size,
        shuffle=True,
        num_workers=4,
        collate_fn=collate_fn,
        pin_memory=True
    )
    
    val_dataset = PretrainDataset(
        data_file=config.data.data_file,
        split="val",
        task="all"
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=config.training.batch_size,
        shuffle=False,
        num_workers=4,
        collate_fn=collate_fn,
        pin_memory=True
    )
    
    print(f"✓ 训练集: {len(train_dataset)} 样本")
    print(f"✓ 验证集: {len(val_dataset)} 样本")
    print("=" * 60)
    
    return train_loader, val_loader


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, required=True,
                       choices=["qwen2.5-vl-3b", "qwen2.5-vl-7b"],
                       help="选择模型")
    parser.add_argument("--epochs", type=int, default=5)
    parser.add_argument("--batch_size", type=int, default=None)
    parser.add_argument("--lr", type=float, default=None)
    
    args = parser.parse_args()
    
    # 选择配置
    if args.model == "qwen2.5-vl-3b":
        config = QWEN25_VL_3B_CONFIG
    elif args.model == "qwen2.5-vl-7b":
        config = QWEN25_VL_7B_CONFIG
    
    # 覆盖配置
    if args.epochs:
        config.training.num_epochs = args.epochs
    if args.batch_size:
        config.training.batch_size = args.batch_size
    if args.lr:
        config.training.learning_rate = args.lr
    
    # 设置随机种子
    set_seed(config.training.seed)
    
    # 打印配置
    print("=" * 60)
    print("配置信息")
    print("=" * 60)
    print(f"模型: {config.model.model_name}")
    print(f"输出: {config.training.output_dir}")
    print(f"Epochs: {config.training.num_epochs}")
    print(f"Batch: {config.training.batch_size}")
    print(f"LR: {config.training.learning_rate}")
    print("=" * 60)
    
    # 创建数据加载器
    train_loader, val_loader = create_dataloaders(config)
    
    # 创建训练器
    trainer = MultiTaskTrainer(config, train_loader, val_loader)
    
    # 开始训练
    trainer.train()
    
    print(f"\n✅ 完成！模型保存在: {config.training.output_dir}")


if __name__ == "__main__":
    main()