"""
训练数据过滤可视化脚本

用途：从训练集中筛选符合条件的样本并进行可视化

功能：
- 支持类别过滤（如：只可视化包含两轮车的样本）
- 支持空间范围过滤（如：只可视化指定3D空间范围内的目标）
- 支持批量处理

使用方法：
    # 可视化包含两轮车的样本（两轮车在左右10米、纵向50米范围内）
    python scripts_for_gt/visualize_filtered_train.py \
        --data data/mono3d.yaml \
        --target-class bicycle \
        --x-range -10 10 \
        --z-range 0 50 \
        --max-samples 20

    # 可视化包含行人的样本
    python scripts_for_gt/visualize_filtered_train.py \
        --data data/mono3d.yaml \
        --target-class pedestrian \
        --max-samples 50

    # 可视化特定空间范围内的所有车辆
    python scripts_for_gt/visualize_filtered_train.py \
        --data data/mono3d.yaml \
        --target-class vehicle \
        --x-range -5 5 \
        --z-range 10 30 \
        --max-samples 30
"""

import argparse
import sys
from pathlib import Path
import yaml
import numpy as np

# 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,
    load_label_file,
)


# 类别名称到ID的映射
CLASS_NAME_TO_ID = {
    'vehicle': 0,
    'pedestrian': 1,
    'bicycle': 2,
    'rider': 3,
    'tricycle': 13,
}

CLASS_ID_TO_NAME = {v: k for k, v in CLASS_NAME_TO_ID.items()}


def check_sample_matches_filter(
    label_path,
    target_class_id=None,
    x_range=None,
    y_range=None,
    z_range=None,
    min_targets=1,
):
    """检查样本是否符合过滤条件
    
    Args:
        label_path: 标签文件路径
        target_class_id: 目标类别ID（None表示不过滤类别）
        x_range: X坐标范围 (min, max)，单位：米，None表示不过滤
        y_range: Y坐标范围 (min, max)，单位：米，None表示不过滤
        z_range: Z坐标范围 (min, max)，单位：米，None表示不过滤
        min_targets: 至少需要多少个符合条件的目标
        
    Returns:
        tuple: (matches, matched_count, matched_objects_info)
            - matches: bool，是否符合条件
            - matched_count: int，符合条件的目标数量
            - matched_objects_info: list of dict，符合条件的目标信息
    """
    try:
        labels = load_label_file(label_path)
    except Exception as e:
        print(f"  警告：无法读取标签文件 {label_path}: {e}")
        return False, 0, []
    
    if len(labels) == 0:
        return False, 0, []
    
    matched_objects = []
    
    for label in labels:
        # 检查类别
        cls = int(label[0])
        if target_class_id is not None and cls != target_class_id:
            continue
        
        # 检查是否有3D信息
        x3d, y3d, z3d = label[6:9]
        if np.isnan(x3d) or np.isnan(y3d) or np.isnan(z3d):
            continue
        
        # 检查空间范围
        in_range = True
        
        if x_range is not None:
            if not (x_range[0] <= x3d <= x_range[1]):
                in_range = False
        
        if y_range is not None:
            if not (y_range[0] <= y3d <= y_range[1]):
                in_range = False
        
        if z_range is not None:
            if not (z_range[0] <= z3d <= z_range[1]):
                in_range = False
        
        if in_range:
            obj_info = {
                'class': cls,
                'class_name': CLASS_ID_TO_NAME.get(cls, f'class_{cls}'),
                'x3d': x3d,
                'y3d': y3d,
                'z3d': z3d,
                'bbox_2d': label[1:5],
            }
            matched_objects.append(obj_info)
    
    matched_count = len(matched_objects)
    matches = matched_count >= min_targets
    
    return matches, matched_count, matched_objects


def find_filtered_samples(
    data_yaml_path,
    target_class=None,
    x_range=None,
    y_range=None,
    z_range=None,
    min_targets=1,
    max_samples=None,
    split='train',
):
    """从数据集中查找符合过滤条件的样本
    
    Args:
        data_yaml_path: 数据集YAML配置文件路径
        target_class: 目标类别名称（如'bicycle'）
        x_range: X坐标范围 (min, max)
        y_range: Y坐标范围 (min, max)
        z_range: Z坐标范围 (min, max)
        min_targets: 最少目标数量
        max_samples: 最大返回样本数
        split: 数据集划分('train'或'val')
        
    Returns:
        list: 符合条件的样本信息列表
    """
    # 读取数据集配置
    with open(data_yaml_path, 'r') as f:
        data = yaml.safe_load(f)
    
    # 获取数据集路径
    dataset_path = data.get(split)
    if dataset_path is None:
        print(f"错误：数据集配置中未找到'{split}'字段")
        return []
    
    # 如果是相对路径，相对于yaml文件的位置
    dataset_path = Path(dataset_path)
    if not dataset_path.is_absolute():
        dataset_path = Path(data_yaml_path).parent / dataset_path
    
    # 读取图像列表
    if dataset_path.is_file():
        # 如果是文件列表
        with open(dataset_path, 'r') as f:
            image_files = [line.strip() for line in f.readlines() if line.strip()]
    else:
        # 如果是目录
        image_files = list(dataset_path.glob("**/*.jpg"))
        image_files.extend(list(dataset_path.glob("**/*.png")))
    
    print(f"在{split}集中找到 {len(image_files)} 个图像文件")
    
    # 转换类别名称到ID
    target_class_id = None
    if target_class is not None:
        if target_class.lower() in CLASS_NAME_TO_ID:
            target_class_id = CLASS_NAME_TO_ID[target_class.lower()]
        else:
            print(f"警告：未知类别 '{target_class}'，将不过滤类别")
    
    # 打印过滤条件
    print("\n过滤条件:")
    if target_class_id is not None:
        print(f"  目标类别: {target_class} (ID={target_class_id})")
    else:
        print(f"  目标类别: 所有类别")
    
    if x_range is not None:
        print(f"  X范围: [{x_range[0]:.1f}, {x_range[1]:.1f}] 米 (左右)")
    if y_range is not None:
        print(f"  Y范围: [{y_range[0]:.1f}, {y_range[1]:.1f}] 米 (上下)")
    if z_range is not None:
        print(f"  Z范围: [{z_range[0]:.1f}, {z_range[1]:.1f}] 米 (前后)")
    print(f"  最少目标数: {min_targets}")
    
    # 遍历查找符合条件的样本
    matched_samples = []
    
    for i, image_file in enumerate(image_files):
        if (i + 1) % 100 == 0:
            print(f"  进度: {i+1}/{len(image_files)}, 已找到 {len(matched_samples)} 个符合条件的样本")
        
        image_rel_path = Path(image_file)
        image_path = dataset_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'))
        
        if not label_path.exists():
            continue
        
        # 检查是否符合过滤条件
        matches, matched_count, matched_objects = check_sample_matches_filter(
            label_path,
            target_class_id=target_class_id,
            x_range=x_range,
            y_range=y_range,
            z_range=z_range,
            min_targets=min_targets,
        )
        
        if matches:
            # 推断标定路径
            calib_path = image_path.parent.parent / 'calib' / 'L2_calib' / 'camera4.json'
            
            sample_info = {
                'image_path': str(image_path),
                'label_path': str(label_path),
                'calib_path': str(calib_path) if calib_path.exists() else None,
                'matched_count': matched_count,
                'matched_objects': matched_objects,
            }
            matched_samples.append(sample_info)
            
            # 检查是否达到最大样本数
            if max_samples is not None and len(matched_samples) >= max_samples:
                print(f"\n已找到 {len(matched_samples)} 个符合条件的样本（达到最大限制）")
                break
    
    print(f"\n总共找到 {len(matched_samples)} 个符合条件的样本")
    
    return matched_samples


def visualize_filtered_samples(
    data_yaml_path,
    output_dir="./filtered_train_viz",
    target_class=None,
    x_range=None,
    y_range=None,
    z_range=None,
    min_targets=1,
    max_samples=20,
    use_roi=False,
    split='train',
    save_types=None,
):
    """可视化符合条件的样本
    
    Args:
        data_yaml_path: 数据集YAML配置文件路径
        output_dir: 输出目录
        target_class: 目标类别名称
        x_range: X坐标范围
        y_range: Y坐标范围
        z_range: Z坐标范围
        min_targets: 最少目标数量
        max_samples: 最大样本数
        use_roi: 是否使用ROI变换
        split: 数据集划分
        save_types: 保存类型列表 ['2d', '3d', 'bev', 'combined']，None表示全部保存
    """
    # 读取数据集配置
    with open(data_yaml_path, 'r') as f:
        data = yaml.safe_load(f)
    
    # 查找符合条件的样本
    matched_samples = find_filtered_samples(
        data_yaml_path=data_yaml_path,
        target_class=target_class,
        x_range=x_range,
        y_range=y_range,
        z_range=z_range,
        min_targets=min_targets,
        max_samples=max_samples,
        split=split,
    )
    
    if len(matched_samples) == 0:
        print("\n未找到符合条件的样本！")
        return
    
    # 获取配置参数
    roi = data.get('roi') if use_roi else None
    virtual_fx = data.get('virtual_fx') if use_roi else None
    ori_img_size = data.get('ori_img_size') if use_roi else None
    
    # 类别名称
    names = {
        0: "vehicle",
        1: "pedestrian",
        2: "bicycle",
        3: "rider",
        4: "motorcycle",
        13: "tricycle",
    }
    
    print(f"\n开始可视化 {len(matched_samples)} 个样本...")
    
    success_count = 0
    error_count = 0
    
    for i, sample in enumerate(matched_samples):
        print(f"\n[{i+1}/{len(matched_samples)}] 处理样本:")
        print(f"  图像: {Path(sample['image_path']).name}")
        print(f"  符合条件的目标数: {sample['matched_count']}")
        
        # 打印目标信息
        for j, obj in enumerate(sample['matched_objects']):
            print(f"    目标{j+1}: {obj['class_name']}, "
                  f"位置=({obj['x3d']:.2f}, {obj['y3d']:.2f}, {obj['z3d']:.2f})m")
        
        try:
            # 创建子目录
            sample_name = Path(sample['image_path']).stem
            sample_output_dir = Path(output_dir)
            
            # 可视化
            visualize_single_frame(
                image_path=sample['image_path'],
                label_path=sample['label_path'],
                output_dir=str(sample_output_dir),
                calib_path=sample['calib_path'],
                roi_size=tuple(roi) if roi else None,
                virtual_fx=virtual_fx,
                ori_img_size=tuple(ori_img_size) if ori_img_size else None,
                names=names,
                save_types=save_types,
            )
            success_count += 1
        except Exception as e:
            print(f"  错误: {e}")
            error_count += 1
    
    print(f"\n===== 可视化完成 =====")
    print(f"成功: {success_count}")
    print(f"失败: {error_count}")
    print(f"输出目录: {output_dir}")
    
    # 生成摘要文件
    summary_path = Path(output_dir) / "summary.txt"
    with open(summary_path, 'w', encoding='utf-8') as f:
        f.write("过滤可视化摘要\n")
        f.write("=" * 50 + "\n\n")
        
        f.write("过滤条件:\n")
        if target_class:
            f.write(f"  目标类别: {target_class}\n")
        if x_range:
            f.write(f"  X范围: [{x_range[0]:.1f}, {x_range[1]:.1f}] 米\n")
        if y_range:
            f.write(f"  Y范围: [{y_range[0]:.1f}, {y_range[1]:.1f}] 米\n")
        if z_range:
            f.write(f"  Z范围: [{z_range[0]:.1f}, {z_range[1]:.1f}] 米\n")
        f.write(f"  最少目标数: {min_targets}\n")
        f.write(f"  数据集: {split}\n\n")
        
        f.write(f"可视化结果:\n")
        f.write(f"  总样本数: {len(matched_samples)}\n")
        f.write(f"  成功: {success_count}\n")
        f.write(f"  失败: {error_count}\n\n")
        
        f.write("样本列表:\n")
        f.write("-" * 50 + "\n")
        for i, sample in enumerate(matched_samples):
            f.write(f"{i+1}. {Path(sample['image_path']).name}\n")
            f.write(f"   符合条件的目标数: {sample['matched_count']}\n")
            for obj in sample['matched_objects']:
                f.write(f"   - {obj['class_name']}: ({obj['x3d']:.2f}, {obj['y3d']:.2f}, {obj['z3d']:.2f})m\n")
            f.write("\n")
    
    print(f"摘要已保存: {summary_path}")


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="/data1/dongying/Mono3d/G1M3/filtered_train_viz",
                       help="输出目录")
    parser.add_argument("--split", type=str, default="train", choices=['train', 'val'],
                       help="数据集划分（train或val）")
    
    # 过滤条件
    parser.add_argument("--target-class", type=str, default='bicycle',
                       choices=['vehicle','pedestrian','bicycle'],
                       help="目标类别")
    parser.add_argument("--x-range", type=float, nargs=2, default=[-10, 10], metavar=("MIN", "MAX"),
                       help="X坐标范围（左右），单位：米，如: --x-range -10 10")
    parser.add_argument("--y-range", type=float, nargs=2, default=None, metavar=("MIN", "MAX"),
                       help="Y坐标范围（上下），单位：米")
    parser.add_argument("--z-range", type=float, nargs=2, default=[0, 50], metavar=("MIN", "MAX"),
                       help="Z坐标范围（前后），单位：米，如: --z-range 0 50")
    parser.add_argument("--min-targets", type=int, default=1,
                       help="最少目标数量")
    parser.add_argument("--max-samples", type=int, default=5000,
                       help="最大可视化样本数")
    
    # 可视化选项
    parser.add_argument("--use-roi", action="store_true",
                       help="使用ROI变换")
    parser.add_argument("--save-types", type=str, nargs='+', 
                       choices=['2d', '3d', 'bev', 'combined'],
                       default=['combined'],
                       help="保存的图像类型（可多选），默认只保存combined")
    
    return parser.parse_args()


def main():
    args = parse_args()
    
    # 验证范围参数
    if args.x_range and args.x_range[0] >= args.x_range[1]:
        print("错误：X范围的最小值必须小于最大值")
        return
    if args.y_range and args.y_range[0] >= args.y_range[1]:
        print("错误：Y范围的最小值必须小于最大值")
        return
    if args.z_range and args.z_range[0] >= args.z_range[1]:
        print("错误：Z范围的最小值必须小于最大值")
        return
    
    visualize_filtered_samples(
        data_yaml_path=args.data,
        output_dir=args.output,
        target_class=args.target_class,
        x_range=tuple(args.x_range) if args.x_range else None,
        y_range=tuple(args.y_range) if args.y_range else None,
        z_range=tuple(args.z_range) if args.z_range else None,
        min_targets=args.min_targets,
        max_samples=args.max_samples,
        use_roi=args.use_roi,
        split=args.split,
        save_types=args.save_types,
    )


if __name__ == "__main__":
    main()
