#!/usr/bin/env python3
"""
测试脚本：验证V2数据准备流程
包括DAD整合和annotation enhancement
"""

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


def test_annotation_enhancement():
    """测试annotation enhancement结果"""
    print("=" * 70)
    print("1. 测试 Annotation Enhancement")
    print("=" * 70)
    
    enhanced_count = 0
    total_count = 0
    
    # 测试DADA-2000
    dada_root = Path("PROJECT_ROOT/data/dataset/pretrain/DADA-2000")
    if dada_root.exists():
        for anno_file in list(dada_root.rglob("annotation.json"))[:10]:  # 抽查10个
            with open(anno_file) as f:
                data = json.load(f)
            
            total_count += 1
            if data.get("annotation_enhanced", False):
                enhanced_count += 1
                print(f"\n✓ 已增强案例:")
                print(f"  原始: {data.get('accident_type_original', 'N/A')}")
                print(f"  增强: {data.get('accident_type', 'N/A')[:80]}...")
    
    if total_count > 0:
        print(f"\n抽查结果: {enhanced_count}/{total_count} 被增强 ({enhanced_count/total_count*100:.1f}%)")
    else:
        print("\n⚠️  未找到DADA-2000数据或annotation未增强")
    
    return total_count > 0


def test_dad_extraction():
    """测试DAD帧提取"""
    print("\n" + "=" * 70)
    print("2. 测试 DAD数据提取")
    print("=" * 70)
    
    dad_frames_dir = Path("PROJECT_ROOT/data/dataset/pretrain/dad_frames")
    
    if not dad_frames_dir.exists():
        print("❌ DAD frames目录不存在")
        print("请运行: python prepare_pretrain_data_v2.py")
        return False
    
    # 统计DAD cases
    dad_cases = list(dad_frames_dir.iterdir())
    positive = [c for c in dad_cases if "positive" in c.name]
    negative = [c for c in dad_cases if "negative" in c.name]
    
    print(f"\n✓ DAD cases总数: {len(dad_cases)}")
    print(f"  Positive: {len(positive)}")
    print(f"  Negative: {len(negative)}")
    
    # 检查几个case的帧数
    if len(dad_cases) > 0:
        sample_case = dad_cases[0]
        frames = list(sample_case.glob("*.jpg"))
        anno_file = sample_case / "annotation.json"
        
        print(f"\n示例case: {sample_case.name}")
        print(f"  帧数: {len(frames)}")
        
        if anno_file.exists():
            with open(anno_file) as f:
                anno = json.load(f)
            print(f"  标注:")
            print(f"    accident: {anno.get('accident')}")
            print(f"    accident_type: {anno.get('accident_type')}")
            print(f"    fps: {anno.get('fps')}")
        else:
            print("  ⚠️  annotation.json不存在")
    
    return len(dad_cases) > 0


def test_data_preparation():
    """测试完整数据准备"""
    print("\n" + "=" * 70)
    print("3. 测试 完整数据准备")
    print("=" * 70)
    
    data_file = Path("PROJECT_ROOT/data/dataset/pretrain/train/pretrain_data_v2.pkl")
    summary_file = Path("PROJECT_ROOT/data/dataset/pretrain/train/pretrain_summary_v2.json")
    
    # 检查文件
    if not data_file.exists():
        print(f"❌ 数据文件不存在: {data_file}")
        print("请运行: python prepare_pretrain_data_v2.py")
        return False
    
    print(f"✓ 数据文件: {data_file}")
    
    # 加载数据
    with open(data_file, "rb") as f:
        data = pickle.load(f)
    
    # 验证结构
    required_splits = ["train", "val", "test"]
    for split in required_splits:
        if split not in data:
            print(f"❌ 缺少split: {split}")
            return False
    
    print("✓ 数据结构正确")
    
    # 统计
    print("\n" + "-" * 70)
    print("数据统计:")
    print("-" * 70)
    
    for split in required_splits:
        split_data = data[split]
        
        n_cases = split_data.get("total_cases", 0)
        n_task1 = len(split_data.get("task1_scene_understanding", []))
        n_task2 = len(split_data.get("task2_binary_detection", []))
        n_task3 = len(split_data.get("task3_accident_description", []))
        n_task4 = len(split_data.get("task4_sequence_prediction", []))
        
        print(f"\n{split.upper()}: {n_cases} cases")
        print(f"  任务1 (场景理解):   {n_task1}")
        print(f"  任务2 (二分类):     {n_task2}")
        print(f"  任务3 (事故描述):   {n_task3}")
        print(f"  任务4 (序列预测):   {n_task4}")
        print(f"  ─────────────────────")
        print(f"  总样本:            {n_task1 + n_task2 + n_task3 + n_task4}")
    
    # 数据集来源分布
    print("\n" + "-" * 70)
    print("数据集来源分布 (训练集):")
    print("-" * 70)
    
    dataset_counts = Counter()
    for task_name in ["task1_scene_understanding", "task2_binary_detection", 
                      "task3_accident_description", "task4_sequence_prediction"]:
        for sample in data["train"].get(task_name, []):
            dataset_counts[sample["metadata"]["dataset"]] += 1
    
    for dataset, count in dataset_counts.items():
        print(f"  {dataset}: {count} 样本")
    
    # DAD占比
    if "dad" in dataset_counts:
        total = sum(dataset_counts.values())
        dad_ratio = dataset_counts["dad"] / total * 100
        print(f"\n✓ DAD数据占比: {dad_ratio:.1f}%")
    else:
        print("\n⚠️  未检测到DAD数据")
    
    # 难度分布
    print("\n" + "-" * 70)
    print("难度分布 (训练集):")
    print("-" * 70)
    
    difficulty_counts = Counter()
    for task_name in ["task1_scene_understanding", "task2_binary_detection",
                      "task3_accident_description", "task4_sequence_prediction"]:
        for sample in data["train"].get(task_name, []):
            difficulty_counts[sample.get("difficulty", "unknown")] += 1
    
    for diff, count in difficulty_counts.items():
        total = sum(difficulty_counts.values())
        print(f"  {diff}: {count} ({count/total*100:.1f}%)")
    
    # 样本示例
    print("\n" + "-" * 70)
    print("样本示例:")
    print("-" * 70)
    
    # Task 1
    task1_samples = data["train"]["task1_scene_understanding"][:2]
    print("\n任务1 - 场景理解:")
    for i, s in enumerate(task1_samples, 1):
        print(f"  样本{i}:")
        print(f"    难度: {s['difficulty']}")
        print(f"    标签: {s['label']}")
    
    # Task 2
    task2_samples = data["train"]["task2_binary_detection"][:2]
    print("\n任务2 - 二分类:")
    for i, s in enumerate(task2_samples, 1):
        print(f"  样本{i}:")
        print(f"    难度: {s['difficulty']}")
        print(f"    标签: {s['label']}")
        print(f"    正样本: {s['metadata']['is_positive']}")
    
    # Task 3
    task3_samples = data["train"]["task3_accident_description"][:1]
    if task3_samples:
        print("\n任务3 - 事故描述:")
        s = task3_samples[0]
        print(f"  难度: {s['difficulty']}")
        print(f"  标签: {s['label'][:80]}...")
        print(f"  增强标注: {s['metadata'].get('was_enhanced', False)}")
    
    # Task 4
    task4_samples = data["train"]["task4_sequence_prediction"][:1]
    if task4_samples:
        print("\n任务4 - 序列预测:")
        s = task4_samples[0]
        print(f"  难度: {s['difficulty']}")
        print(f"  序列长度: {s['metadata']['sequence_length']}")
        print(f"  标签: {s['label'][:80]}...")
    
    return True


def main():
    """主测试流程"""
    print("\n" + "=" * 70)
    print("LKAlert预训练数据准备 - 测试脚本 V2")
    print("=" * 70)
    
    results = {}
    
    # Test 1: Annotation Enhancement
    results["annotation"] = test_annotation_enhancement()
    
    # Test 2: DAD Extraction
    results["dad"] = test_dad_extraction()
    
    # Test 3: Data Preparation
    results["data_prep"] = test_data_preparation()
    
    # 总结
    print("\n" + "=" * 70)
    print("测试总结")
    print("=" * 70)
    
    all_passed = all(results.values())
    
    if all_passed:
        print("✅ 所有测试通过!")
        print("\n准备就绪，可以开始训练:")
        print("  bash run_pretrain_v2.sh qwen2.5-vl-3b")
    else:
        print("❌ 部分测试失败:")
        for test_name, passed in results.items():
            status = "✓" if passed else "✗"
            print(f"  {status} {test_name}")
        
        print("\n请按照提示修复问题")
    
    return all_passed


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)