import os
import random
import shutil
from collections import defaultdict

def parse_sysu_mm01_structure(data_root):
    """
    解析SYSU-MM01数据集结构，按相机和ID分类
    :param data_root: 数据集根目录
    :return: 结构为 {相机号: {ID: [图像路径列表]}} 的字典
    """
    # 定义相机类型和场景属性
    cam_info = {
        1: {'type': 'rgb', 'scene': 'indoor'},
        2: {'type': 'rgb', 'scene': 'indoor'},
        3: {'type': 'ir', 'scene': 'indoor'},
        4: {'type': 'rgb', 'scene': 'outdoor'},
        5: {'type': 'rgb', 'scene': 'outdoor'},
        6: {'type': 'ir', 'scene': 'outdoor'}
    }
    
    # 初始化数据结构
    cam_dict = defaultdict(lambda: defaultdict(list))  # cam -> id -> [img_paths]
    
    # 遍历数据集目录
    for cam in cam_info.keys():
        cam_dir = os.path.join(data_root, f'cam{cam}')
        if not os.path.exists(cam_dir):
            print(f"警告: 相机目录 {cam_dir} 不存在")
            continue
            
        # 遍历每个ID目录
        for id_folder in os.listdir(cam_dir):
            id_path = os.path.join(cam_dir, id_folder)
            if not os.path.isdir(id_path):
                continue
                
            # 收集该ID下的所有图像
            for img_file in os.listdir(id_path):
                if img_file.endswith(('.jpg', '.png')):  # 假设图像格式为jpg或png
                    img_path = os.path.join(id_path, img_file)
                    cam_dict[cam][id_folder].append(img_path)
    
    return cam_dict, cam_info

def select_samples(cam_dict, cam_info, output_dir):
    """
    按照指定策略抽取样本
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'rgb'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'ir'), exist_ok=True)
    
    # 记录抽取信息
    selection_log = []
    
    # 处理RGB相机 (1,2,4,5) - 每个相机抽取50个ID，每个ID1张
    rgb_cams = [1, 2, 4, 5]
    for cam in rgb_cams:
        cam_id_dict = cam_dict.get(cam, {})
        if not cam_id_dict:
            print(f"警告: 相机 {cam} 没有找到数据")
            continue
            
        # 随机选择50个ID (如果ID数量不足50则全部选择)
        selected_ids = random.sample(list(cam_id_dict.keys()), min(50, len(cam_id_dict)))
        
        for idx, person_id in enumerate(selected_ids):
            # 从该ID中随机选择一张图像
            img_paths = cam_id_dict[person_id]
            selected_img = random.choice(img_paths)
            
            # 复制图像到输出目录
            dest_path = os.path.join(output_dir, 'rgb', 
                                   f'cam{cam}_id{person_id}_{os.path.basename(selected_img)}')
            shutil.copy2(selected_img, dest_path)
            
            # 记录选择信息
            selection_log.append({
                'type': 'rgb',
                'cam': cam,
                'scene': cam_info[cam]['scene'],
                'person_id': person_id,
                'src': selected_img,
                'dest': dest_path
            })
            
            print(f"已抽取 RGB 图像: 相机{cam} ID{person_id} ({idx+1}/50)")
    
    # 处理IR相机 (3,6) - 每个相机抽取100个ID，每个ID1张
    ir_cams = [3, 6]
    for cam in ir_cams:
        cam_id_dict = cam_dict.get(cam, {})
        if not cam_id_dict:
            print(f"警告: 相机 {cam} 没有找到数据")
            continue
            
        # 随机选择100个ID (如果ID数量不足100则全部选择)
        selected_ids = random.sample(list(cam_id_dict.keys()), min(100, len(cam_id_dict)))
        
        for idx, person_id in enumerate(selected_ids):
            # 从该ID中随机选择一张图像
            img_paths = cam_id_dict[person_id]
            selected_img = random.choice(img_paths)
            
            # 复制图像到输出目录
            dest_path = os.path.join(output_dir, 'ir', 
                                   f'cam{cam}_id{person_id}_{os.path.basename(selected_img)}')
            shutil.copy2(selected_img, dest_path)
            
            # 记录选择信息
            selection_log.append({
                'type': 'ir',
                'cam': cam,
                'scene': cam_info[cam]['scene'],
                'person_id': person_id,
                'src': selected_img,
                'dest': dest_path
            })
            
            print(f"已抽取 IR 图像: 相机{cam} ID{person_id} ({idx+1}/100)")
    
    # 保存抽取日志
    with open(os.path.join(output_dir, 'selection_log.txt'), 'w') as f:
        for entry in selection_log:
            f.write(f"{entry['type']} | 相机{entry['cam']}({entry['scene']}) | ID{entry['person_id']} | 来源: {entry['src']}\n")
    
    print(f"\n抽取完成! 共抽取 {len(selection_log)} 张图像")
    print(f"输出目录: {output_dir}")
    print(f"抽取日志: {os.path.join(output_dir, 'selection_log.txt')}")

if __name__ == "__main__":
    # 配置路径
    DATA_ROOT = "/path/to/sysu-mm01"  # 替换为你的SYSU-MM01数据集根目录
    OUTPUT_DIR = "./selected_samples"  # 输出目录
    
    # 解析数据集结构
    print("正在解析数据集结构...")
    cam_dict, cam_info = parse_sysu_mm01_structure(DATA_ROOT)
    
    # 抽取样本
    print("\n开始抽取样本...")
    select_samples(cam_dict, cam_info, OUTPUT_DIR)