#!/opt/homebrew/bin/python3.11
"""
YOLO Detection Count Comparison
对比 YOLO processor 的检测物体数量和类别
"""

import json
from pathlib import Path
from collections import defaultdict

def load_yolo_results(filepath):
    """加载 YOLO 检测结果"""
    if not filepath.exists():
        print(f"❌ 文件不存在: {filepath}")
        return None
    
    with open(filepath) as f:
        data = json.load(f)
    
    results = {
        'total_frames': 0,
        'total_objects': 0,
        'classes': defaultdict(int),
        'confidence_avg': 0.0,
        'frames_with_objects': 0,
        'frames_empty': 0,
    }
    
    frames_data = data.get('frames', [])
    
    if isinstance(frames_data, dict):
        frames_list = list(frames_data.values())
    elif isinstance(frames_data, list):
        frames_list = frames_data
    else:
        return results
    
    results['total_frames'] = len(frames_list)
    
    total_confidence = 0.0
    total_objects = 0
    
    for frame in frames_list:
        if isinstance(frame, dict):
            detections = frame.get('detections', frame.get('objects', []))
            
            if detections and len(detections) > 0:
                results['frames_with_objects'] += 1
                
                for det in detections:
                    if isinstance(det, dict):
                        total_objects += 1
                        class_name = det.get('class_name', det.get('name', det.get('class', 'unknown')))
                        results['classes'][class_name] += 1
                        
                        confidence = det.get('confidence', det.get('score', 1.0))
                        if isinstance(confidence, (int, float)):
                            total_confidence += confidence
            else:
                results['frames_empty'] += 1
    
    results['total_objects'] = total_objects
    if total_objects > 0:
        results['confidence_avg'] = total_confidence / total_objects
    
    return results

def compare_classes(results_a, results_b, results_c):
    """对比检测类别"""
    all_classes = set()
    
    if results_a:
        all_classes.update(results_a['classes'].keys())
    if results_b:
        all_classes.update(results_b['classes'].keys())
    if results_c:
        all_classes.update(results_c['classes'].keys())
    
    comparison = []
    
    for class_name in sorted(all_classes):
        count_a = results_a['classes'].get(class_name, 0) if results_a else 0
        count_b = results_b['classes'].get(class_name, 0) if results_b else 0
        count_c = results_c['classes'].get(class_name, 0) if results_c else 0
        
        if count_a != count_b or count_a != count_c or count_b != count_c:
            comparison.append({
                'class_name': class_name,
                'cpu': count_a,
                'mps': count_b,
                'contract': count_c,
                'max': max(count_a, count_b, count_c),
                'min': min(count_a, count_b, count_c),
            })
    
    return comparison

def main():
    benchmark_dir = Path('/Users/accusys/momentry_core_0.1/output/benchmark/yolo_processor')
    
    print("=" * 80)
    print("YOLO Detection Count Comparison")
    print("=" * 80)
    print()
    
    # 加载三个版本的结果
    results_a = load_yolo_results(benchmark_dir / 'scheme_A_yolo_processor.json')
    results_b = load_yolo_results(benchmark_dir / 'scheme_B_yolo_processor_mps.json')
    results_c = load_yolo_results(benchmark_dir / 'scheme_C_yolo_processor_contract_v1.json')
    
    if not results_a and not results_b and not results_c:
        print("❌ 没有可用的检测结果文件")
        return
    
    # 统计概览
    print("【检测统计】")
    print()
    
    print("| 版本 | 总帧数 | 检测物体数 | 有物体帧 | 空帧 | 平均置信度 |")
    print("|------|--------|-----------|---------|------|------------|")
    
    for name, results in [('CPU', results_a), ('MPS', results_b), ('Contract', results_c)]:
        if results:
            print(f"| {name} | {results['total_frames']} | {results['total_objects']} | {results['frames_with_objects']} | {results['frames_empty']} | {results['confidence_avg']:.2f} |")
        else:
            print(f"| {name} | - | - | - | - | - |")
    
    print()
    
    # 类别统计
    print("【检测类别统计】")
    print()
    
    if results_a and results_a['classes']:
        print("CPU版本检测类别:")
        print("| 类别 | 数量 |")
        print("|------|------|")
        
        for class_name, count in sorted(results_a['classes'].items(), key=lambda x: -x[1]):
            print(f"| {class_name} | {count} |")
        
        print(f"| **总计** | {sum(results_a['classes'].values())} |")
    
    print()
    
    # 类别对比
    print("【类别数量对比】")
    print()
    
    comparison = compare_classes(results_a, results_b, results_c)
    
    if comparison:
        print(f"共有 {len(comparison)} 个类别检测数量不同")
        print()
        
        print("| 类别 | CPU | MPS | Contract | 最大差异 |")
        print("|------|-----|-----|----------|---------|")
        
        for item in comparison[:20]:
            diff = item['max'] - item['min']
            print(f"| {item['class_name']} | {item['cpu']} | {item['mps']} | {item['contract']} | {diff} |")
        
        if len(comparison) > 20:
            print("| ... | ... | ... | ... | ... |")
    else:
        print("所有类别检测数量一致")
    
    print()
    
    # 检测率分析
    print("【检测率分析】")
    print()
    
    if results_a and results_a['total_objects'] > 0:
        baseline = results_a['total_objects']
        print(f"以CPU版本为基准（{baseline}个物体）:")
        print()
        
        print("| 版本 | 检测数 | 检测率 | 漏检数 |")
        print("|------|--------|--------|--------|")
        
        for name, results in [('CPU', results_a), ('MPS', results_b), ('Contract', results_c)]:
            if results:
                rate = results['total_objects'] / baseline * 100
                missed = baseline - results['total_objects']
                print(f"| {name} | {results['total_objects']} | {rate:.1f}% | {missed} |")
            else:
                print(f"| {name} | - | - | - |")
    
    print()
    print("=" * 80)
    print("对比完成")
    print("=" * 80)
    
    # 保存结果
    output = {
        'cpu': results_a,
        'mps': results_b,
        'contract': results_c,
        'comparison': comparison,
    }
    
    output_path = benchmark_dir / 'YOLO_COUNT_COMPARISON.json'
    with open(output_path, 'w') as f:
        json.dump(output, f, indent=2, ensure_ascii=False)
    
    print(f"\n对比结果已保存: {output_path}")

if __name__ == '__main__':
    main()