"""
快速测试脚本 - 自动查找并可视化第一个有效样本

用途：自动从验证集中查找一个带标签的样本并进行可视化，用于测试脚本功能

使用方法：
    python scripts_for_gt/test_visualization.py --data data/mono3d.yaml
"""

import argparse
import sys
from pathlib import Path
import yaml

# Add project root to path
FILE = Path(__file__).resolve()
ROOT = FILE.parents[1]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))

from visualize_single_frame import visualize_single_frame


def find_first_valid_sample(data_yaml_path):
    """从数据集配置中查找第一个有效样本
    
    Args:
        data_yaml_path: 数据集YAML配置文件路径
        
    Returns:
        tuple: (image_path, label_path, calib_path) 或 None
    """
    # 读取数据集配置
    with open(data_yaml_path, 'r') as f:
        data = yaml.safe_load(f)
    
    # 获取验证集路径
    val_path = data.get('val')
    if val_path is None:
        print("错误：数据集配置中未找到'val'字段")
        return None
    
    # 如果是相对路径，相对于yaml文件的位置
    val_path = Path(val_path)
    if not val_path.is_absolute():
        val_path = Path(data_yaml_path).parent / val_path
    
    # 读取验证集图像列表
    if val_path.is_file():
        # 如果是文件列表
        with open(val_path, 'r') as f:
            image_files = [line.strip() for line in f.readlines() if line.strip()]
    else:
        # 如果是目录
        image_files = list(val_path.glob("**/*.jpg"))
        image_files.extend(list(val_path.glob("**/*.png")))
    
    print(f"在验证集中找到 {len(image_files)} 个图像文件")
    
    # 遍历查找第一个有标签的样本
    for image_path in image_files:
        image_rel_path = Path(image_path)
        
        image_path = val_path.parent / image_rel_path if not image_rel_path.is_absolute() else image_rel_path
        # 推断标签路径
        label_path = Path(str(image_path).replace('/images/', '/labels/').replace('.jpg', '.txt').replace('.png', '.txt'))
        
        # 推断标定路径
        calib_path = image_path.parent.parent / 'calib' / 'L2_calib' / 'camera4.json'
        
        if label_path.exists():
            print(f"\n找到有效样本:")
            print(f"  图像: {image_path}")
            print(f"  标签: {label_path}")
            print(f"  标定: {calib_path if calib_path.exists() else '未找到'}")
            return str(image_path), str(label_path), str(calib_path) if calib_path.exists() else None
    
    print("错误：未找到有效样本")
    return None


def test_visualization(data_yaml_path, output_dir="./test_gt_viz", use_roi=True):
    """测试可视化功能
    
    Args:
        data_yaml_path: 数据集YAML配置文件路径
        output_dir: 输出目录
        use_roi: 是否使用ROI变换
    """
    # 读取数据集配置
    with open(data_yaml_path, 'r') as f:
        data = yaml.safe_load(f)
    
    # 查找第一个有效样本
    result = find_first_valid_sample(data_yaml_path)
    if result is None:
        return
    
    image_path, label_path, calib_path = result
    
    # 获取配置参数
    roi = data.get('roi')  # [width, height]
    virtual_fx = data.get('virtual_fx')
    ori_img_size = data.get('ori_img_size')  # [width, height]
    
    print(f"\n数据集配置:")
    print(f"  ROI: {roi}")
    print(f"  虚拟焦距: {virtual_fx}")
    print(f"  原始图像尺寸: {ori_img_size}")
    
    # 类别名称
    names = {
        0: "vehicle",
        1: "pedestrian",
        2: "bicycle",
        3: "rider",
        13: "tricycle",
    }
    
    print(f"\n开始可视化...")
    
    try:
        # 测试1：不使用ROI
        print("\n=== 测试1：原始图像可视化（不使用ROI）===")
        visualize_single_frame(
            image_path=image_path,
            label_path=label_path,
            output_dir=f"{output_dir}/no_roi",
            calib_path=calib_path,
            roi_size=None,
            virtual_fx=None,
            ori_img_size=None,
            names=names,
        )
        print("✓ 测试1完成")
        
        # 测试2：使用ROI（如果配置了）
        if use_roi and roi is not None and virtual_fx is not None:
            print("\n=== 测试2：ROI变换后可视化 ===")
            visualize_single_frame(
                image_path=image_path,
                label_path=label_path,
                output_dir=f"{output_dir}/with_roi",
                calib_path=calib_path,
                roi_size=tuple(roi),
                virtual_fx=virtual_fx,
                ori_img_size=tuple(ori_img_size) if ori_img_size else None,
                names=names,
            )
            print("✓ 测试2完成")
        
        print(f"\n✅ 所有测试完成！结果保存在: {output_dir}")
        print("\n生成的文件:")
        output_path = Path(output_dir)
        for file in sorted(output_path.rglob("*.jpg")):
            print(f"  {file.relative_to(output_path)}")
        
    except Exception as e:
        print(f"\n❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()


def parse_args():
    parser = argparse.ArgumentParser(description="快速测试真值可视化功能")
    parser.add_argument("--data", type=str, default="data/mono3d.yaml",
                       help="数据集YAML配置文件")
    parser.add_argument("--output", type=str, default="./test_gt_viz",
                       help="输出目录")
    parser.add_argument("--no-roi", action="store_true",
                       help="不测试ROI变换")
    
    return parser.parse_args()


def main():
    args = parse_args()
    
    test_visualization(
        data_yaml_path=args.data,
        output_dir=args.output,
        use_roi=not args.no_roi,
    )


if __name__ == "__main__":
    main()