"""
统计每个类别的图片数量（包含train/val/test）
"""
import os
from pathlib import Path
from collections import defaultdict


def count_images_by_class(dataset_root: str, output_file: str = None):
    """
    统计每个类别的图片数量（包含train/val/test）
    
    Args:
        dataset_root: 数据集根目录（包含train/val/test目录）
        output_file: 输出文件路径（可选，如果提供则保存到文件）
    """
    dataset_root = Path(dataset_root)
    train_dir = dataset_root / 'train'
    val_dir = dataset_root / 'val'
    test_dir = dataset_root / 'test'
    
    if not train_dir.exists():
        print(f"错误: 训练目录不存在: {train_dir}")
        return
    
    print("="*80)
    print("每个类别的图片数量统计（包含train/val/test）")
    print("="*80)
    
    class_counts = {}
    class_counts_detail = {}  # 详细统计：train/val/test分别的数量
    total_images = 0
    total_train = 0
    total_val = 0
    total_test = 0
    
    # 获取所有类别文件夹（从train目录）
    class_folders = sorted([f for f in train_dir.iterdir() if f.is_dir()])
    
    # 也检查val和test目录中是否有train中没有的类别
    all_class_names = set()
    for folder in class_folders:
        all_class_names.add(folder.name)
    
    if val_dir.exists():
        for folder in val_dir.iterdir():
            if folder.is_dir():
                all_class_names.add(folder.name)
    
    if test_dir.exists():
        for folder in test_dir.iterdir():
            if folder.is_dir():
                all_class_names.add(folder.name)
    
    # 统计每个类别的图片数量
    for class_name in sorted(all_class_names):
        train_count = 0
        val_count = 0
        test_count = 0
        
        # 统计train目录
        image_exts = {'.jpg', '.jpeg', '.png'}
        train_class_dir = train_dir / class_name
        if train_class_dir.exists():
            train_count = sum(
                1 for f in train_class_dir.iterdir()
                if f.is_file() and f.suffix.lower() in image_exts
            )
        
        # 统计val目录
        val_class_dir = val_dir / class_name
        if val_class_dir.exists():
            val_count = sum(
                1 for f in val_class_dir.iterdir()
                if f.is_file() and f.suffix.lower() in image_exts
            )
        
        # 统计test目录
        test_class_dir = test_dir / class_name
        if test_class_dir.exists():
            test_count = sum(
                1 for f in test_class_dir.iterdir()
                if f.is_file() and f.suffix.lower() in image_exts
            )
        
        total = train_count + val_count + test_count
        class_counts[class_name] = total
        class_counts_detail[class_name] = {
            'train': train_count,
            'val': val_count,
            'test': test_count,
            'total': total
        }
        total_images += total
        total_train += train_count
        total_val += val_count
        total_test += test_count
    
    # 按数量排序
    sorted_classes = sorted(class_counts.items(), key=lambda x: x[1], reverse=True)
    
    # 修复文件开头的注释
    if output_file:
        # 确保sorted_classes已定义
        pass
    
    # 打印结果
    print(f"\n总类别数: {len(class_counts)}")
    print(f"总图片数: {total_images} (训练集: {total_train}, 验证集: {total_val}, 测试集: {total_test})")
    print(f"平均每类: {total_images // len(class_counts) if len(class_counts) > 0 else 0} 张\n")
    print("-"*100)
    print(f"{'类别名称':<50} {'总计':<8} {'训练':<8} {'验证':<8} {'测试':<8} {'状态':<10}")
    print("-"*100)
    
    # 统计样本不足的类别（假设阈值是10）
    min_threshold = 10
    insufficient_count = 0
    
    for class_name, count in sorted_classes:
        detail = class_counts_detail[class_name]
        status = "⚠样本不足" if count < min_threshold else "✓"
        if count < min_threshold:
            insufficient_count += 1
        print(f"{class_name:<50} {detail['total']:<8} {detail['train']:<8} {detail['val']:<8} {detail['test']:<8} {status:<10}")
    
    print("-"*100)
    print(f"\n样本不足的类别数（<{min_threshold}张）: {insufficient_count}")
    print(f"样本充足的类别数（>={min_threshold}张）: {len(class_counts) - insufficient_count}")
    
    # 保存到文件（如果指定）
    if output_file:
        output_path = Path(output_file)
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write("="*100 + "\n")
            f.write("每个类别的图片数量统计（包含train/val/test）\n")
            f.write("="*100 + "\n\n")
            f.write(f"总类别数: {len(class_counts)}\n")
            f.write(f"总图片数: {total_images} (训练集: {total_train}, 验证集: {total_val}, 测试集: {total_test})\n")
            f.write(f"平均每类: {total_images // len(class_counts) if len(class_counts) > 0 else 0} 张\n\n")
            f.write("-"*100 + "\n")
            f.write(f"{'类别名称':<50} {'总计':<8} {'训练':<8} {'验证':<8} {'测试':<8} {'状态':<10}\n")
            f.write("-"*100 + "\n")
            
            for class_name, count in sorted_classes:
                detail = class_counts_detail[class_name]
                status = "⚠样本不足" if count < min_threshold else "✓"
                f.write(f"{class_name:<50} {detail['total']:<8} {detail['train']:<8} {detail['val']:<8} {detail['test']:<8} {status:<10}\n")
            
            f.write("-"*100 + "\n")
            f.write(f"\n样本不足的类别数（<{min_threshold}张）: {insufficient_count}\n")
            f.write(f"样本充足的类别数（>={min_threshold}张）: {len(class_counts) - insufficient_count}\n")
        
        print(f"\n✓ 统计结果已保存到: {output_path}")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='统计每个类别的图片数量')
    parser.add_argument('--dataset', type=str, 
                       default='dataset/WholeIngredientRecognition',
                       help='数据集根目录（默认: dataset/WholeIngredientRecognition）')
    parser.add_argument('--output', type=str, default=None,
                       help='输出文件路径（可选，保存统计结果到文件）')
    
    args = parser.parse_args()
    
    # 获取项目根目录
    script_dir = Path(__file__).parent
    project_root = script_dir.parent
    
    # 处理相对路径
    if not os.path.isabs(args.dataset):
        dataset_path = project_root / args.dataset
    else:
        dataset_path = Path(args.dataset)
    
    count_images_by_class(str(dataset_path), args.output)