#!/usr/bin/env python3
"""
测试脚本：验证数据准备流程
"""

import json
import pickle
from pathlib import Path
from collections import Counter

def test_data_preparation():
    """测试数据准备结果"""
    
    print("=" * 60)
    print("测试预训练数据准备")
    print("=" * 60)
    
    data_file = Path("PROJECT_ROOT/data/dataset/pretrain/train/pretrain_data.pkl")
    summary_file = Path("PROJECT_ROOT/data/dataset/pretrain/train/pretrain_summary.json")
    
    # 1. 检查文件存在
    print("\n1. 检查文件...")
    if not data_file.exists():
        print(f"❌ 未找到: {data_file}")
        print("请先运行: python prepare_pretrain_data.py")
        return False
    print(f"✓ 找到数据文件: {data_file}")
    
    if not summary_file.exists():
        print(f"⚠️  未找到统计文件: {summary_file}")
    else:
        print(f"✓ 找到统计文件: {summary_file}")
    
    # 2. 加载数据
    print("\n2. 加载数据...")
    with open(data_file, "rb") as f:
        data = pickle.load(f)
    print(f"✓ 数据加载成功")
    
    # 3. 验证数据结构
    print("\n3. 验证数据结构...")
    required_splits = ["train", "val", "test"]
    for split in required_splits:
        if split not in data:
            print(f"❌ 缺少split: {split}")
            return False
        print(f"✓ {split} split存在")
    
    # 4. 统计信息
    print("\n4. 数据统计:")
    print("-" * 60)
    
    total_stats = {
        "train": {"cases": 0, "task1": 0, "task2": 0, "task3": 0},
        "val": {"cases": 0, "task1": 0, "task2": 0, "task3": 0},
        "test": {"cases": 0, "task1": 0, "task2": 0, "task3": 0}
    }
    
    for split in required_splits:
        split_data = data[split]
        
        n_cases = split_data.get("total_cases", 0)
        n_task1 = len(split_data.get("task1_environment", []))
        n_task2 = len(split_data.get("task2_accident_detection", []))
        n_task3 = len(split_data.get("task3_sequence_prediction", []))
        
        total_stats[split]["cases"] = n_cases
        total_stats[split]["task1"] = n_task1
        total_stats[split]["task2"] = n_task2
        total_stats[split]["task3"] = n_task3
        
        print(f"\n{split.upper()}:")
        print(f"  案例数: {n_cases}")
        print(f"  任务1 (环境描述): {n_task1} 样本")
        print(f"  任务2 (事故检测): {n_task2} 样本")
        print(f"  任务3 (序列预测): {n_task3} 样本")
        print(f"  总样本: {n_task1 + n_task2 + n_task3}")
    
    # 5. 数据质量检查
    print("\n5. 数据质量检查:")
    print("-" * 60)
    
    # 检查train集的样本
    train_task1 = data["train"]["task1_environment"][:5]
    train_task2 = data["train"]["task2_accident_detection"][:5]
    train_task3 = data["train"]["task3_sequence_prediction"][:3]
    
    print("\n任务1样本示例:")
    for i, sample in enumerate(train_task1[:2], 1):
        print(f"  样本{i}:")
        print(f"    图像: {sample['image_path']}")
        print(f"    标签: {sample['label']}")
        print(f"    来源: {sample['metadata']['dataset']}")
    
    print("\n任务2样本示例:")
    for i, sample in enumerate(train_task2[:2], 1):
        print(f"  样本{i}:")
        print(f"    图像: {sample['image_path']}")
        print(f"    标签: {sample['label']}")
    
    print("\n任务3样本示例:")
    for i, sample in enumerate(train_task3[:1], 1):
        print(f"  样本{i}:")
        print(f"    序列长度: {len(sample['image_sequence'])}")
        print(f"    首帧: {sample['image_sequence'][0]}")
        print(f"    标签: {sample['label'][:80]}...")
    
    # 6. 检查图像路径
    print("\n6. 验证图像路径...")
    
    test_paths = []
    if train_task1:
        test_paths.append(train_task1[0]["image_path"])
    if train_task2:
        test_paths.append(train_task2[0]["image_path"])
    if train_task3:
        test_paths.append(train_task3[0]["image_sequence"][0])
    
    all_valid = True
    for path in test_paths:
        if not Path(path).exists():
            print(f"❌ 图像不存在: {path}")
            all_valid = False
    
    if all_valid:
        print(f"✓ 抽查的 {len(test_paths)} 个图像路径有效")
    
    # 7. 任务分布统计
    print("\n7. 任务分布:")
    print("-" * 60)
    
    # Task2标签分布
    task2_labels = [s["label"] for s in data["train"]["task2_accident_detection"]]
    label_counts = Counter(task2_labels)
    print(f"\n任务2标签分布:")
    for label, count in label_counts.items():
        print(f"  {label}: {count} ({count/len(task2_labels)*100:.1f}%)")
    
    # Task3事故比例
    task3_samples = data["train"]["task3_sequence_prediction"]
    accident_count = sum(1 for s in task3_samples if "Accident: Yes" in s["label"])
    print(f"\n任务3事故分布:")
    print(f"  有事故: {accident_count} ({accident_count/len(task3_samples)*100:.1f}%)")
    print(f"  无事故: {len(task3_samples)-accident_count} ({(len(task3_samples)-accident_count)/len(task3_samples)*100:.1f}%)")
    
    # 8. 总结
    print("\n" + "=" * 60)
    print("测试总结:")
    print("=" * 60)
    
    total_samples = sum(
        total_stats["train"]["task1"] + 
        total_stats["train"]["task2"] + 
        total_stats["train"]["task3"]
        for _ in ["train"]
    ) + sum(
        total_stats["val"]["task1"] + 
        total_stats["val"]["task2"] + 
        total_stats["val"]["task3"]
        for _ in ["val"]
    ) + sum(
        total_stats["test"]["task1"] + 
        total_stats["test"]["task2"] + 
        total_stats["test"]["task3"]
        for _ in ["test"]
    )
    
    print(f"✓ 总案例数: {sum(total_stats[s]['cases'] for s in required_splits)}")
    print(f"✓ 总样本数: {total_samples}")
    print(f"✓ 数据准备成功!")
    
    return True


if __name__ == "__main__":
    success = test_data_preparation()
    
    if success:
        print("\n✅ 所有检查通过！可以开始训练。")
        print("\n下一步:")
        print("1. 使用 pretrain_dataset.py 加载数据")
        print("2. 编写VLM微调脚本")
        print("3. 开始预训练")
    else:
        print("\n❌ 检查失败，请修复错误后重试。")