#!/usr/bin/env python3
"""
自适应Prompt预训练主脚本
使用根据annotation长度定制的prompt
"""

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/training/pretrain')

from pretrain_dataset_adaptive import AdaptivePretrainDataset, collate_fn_adaptive
from config import QWEN25_VL_3B_CONFIG, QWEN25_VL_7B_CONFIG
from trainer_v2 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, stage=3):
    """创建数据加载器"""
    print("=" * 70)
    print(f"准备数据 - Curriculum Stage {stage}")
    
    train_dataset = AdaptivePretrainDataset(
        data_file=config.data.data_file,
        split="train",
        tasks=["task1", "task2", "task3", "task4"],
        curriculum_stage=stage,
        use_system_prompt=True
    )
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=config.training.batch_size,
        shuffle=True,
        num_workers=4,
        collate_fn=collate_fn_adaptive,
        pin_memory=True
    )
    
    val_dataset = AdaptivePretrainDataset(
        data_file=config.data.data_file,
        split="val",
        tasks=["task1", "task2", "task3", "task4"],
        curriculum_stage=3,  # 验证集始终使用全部数据
        use_system_prompt=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=config.training.batch_size,
        shuffle=False,
        num_workers=4,
        collate_fn=collate_fn_adaptive,
        pin_memory=True
    )
    
    print(f"✓ 训练集: {len(train_dataset)} 样本")
    print(f"✓ 验证集: {len(val_dataset)} 样本")
    print("=" * 70)
    
    return train_loader, val_loader


def main():
    parser = argparse.ArgumentParser(description="自适应Prompt VLM预训练")
    parser.add_argument("--model", type=str, required=True,
                       choices=["qwen2.5-vl-3b", "qwen2.5-vl-7b"],
                       help="选择模型")
    parser.add_argument("--data_file", type=str, 
                       default="PROJECT_ROOT/data/dataset/pretrain/train/pretrain_data_adaptive.pkl",
                       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)
    parser.add_argument("--wandb", action="store_true",
                       help="启用wandb logging")
    
    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
    
    # 覆盖配置
    config.data.data_file = args.data_file
    
    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
    if args.wandb:
        config.training.use_wandb = True
    
    # 设置随机种子
    set_seed(config.training.seed)
    
    # 打印配置
    print("=" * 70)
    print("配置信息")
    print("=" * 70)
    print(f"模型: {config.model.model_name}")
    print(f"数据: {config.data.data_file}")
    print(f"输出: {config.training.output_dir}")
    print(f"Epochs: {config.training.num_epochs}")
    print(f"Batch Size: {config.training.batch_size}")
    print(f"Learning Rate: {config.training.learning_rate}")
    print(f"WandB: {config.training.use_wandb}")
    print("=" * 70)
    print("\n策略: 自适应Prompt")
    print("  - 短标注 (<20字符) → 简单prompt (识别对象)")
    print("  - 详细标注 (>=20字符) → 详细prompt (完整描述)")
    print("=" * 70)
    
    # 检查数据文件
    if not os.path.exists(args.data_file):
        print(f"\n❌ 数据文件不存在: {args.data_file}")
        print("请先运行:")
        print("  1. python analyze_annotations.py  # 分析标注")
        print("  2. python prepare_pretrain_data_adaptive.py  # 准备数据")
        return
    
    # 创建数据加载器
    train_loader, val_loader = create_dataloaders(config, stage=3)
    
    # 创建训练器
    trainer = MultiTaskTrainer(config, train_loader, val_loader)
    
    # 开始训练
    trainer.train()
    
    print(f"\n✅ 完成！模型保存在: {config.training.output_dir}")
    print("\n下一步:")
    print("1. 评估模型性能")
    print("2. 进行SFT阶段训练")


if __name__ == "__main__":
    main()