import numpy as np
import matplotlib
matplotlib.use('Agg')

class EnergyEfficiencyAnalyzer:
    """
    边缘部署能效分析器
    
    核心指标:
    1. mAP/W - 单位功耗的精度
    2. FPS/W - 单位功耗的帧率
    3. mAP·FPS/W - 综合能效
    """
    
    PLATFORMS = {
        'Jetson Orin NX': {'tops': 100, 'tdp_range': (10, 25)},
        'Jetson Nano': {'tops': 0.5, 'tdp_range': (5, 10)},
        'Raspberry Pi 5': {'tops': 5, 'tdp_range': (3, 8)},
        'Intel NUC 11': {'tops': 8, 'tdp_range': (15, 28)},
        'Orange Pi 5': {'tops': 6, 'tdp_range': (3, 5)},
    }
    
    MODELS = {
        'YOLOv8n': {'params': 3.2, 'flops': 8.7, 'mAP': 37.3},
        'YOLOv8s': {'params': 11.2, 'flops': 28.6, 'mAP': 44.9},
        'YOLOv8m': {'params': 25.9, 'flops': 78.9, 'mAP': 50.2},
        'YOLOv8l': {'params': 43.7, 'flops': 165.2, 'mAP': 52.6},
        'RT-DETR-l': {'params': 32.0, 'flops': 108.0, 'mAP': 53.2},
    }
    
    def estimate_fps(self, model: str, platform: str, quantization: str = 'int8') -> float:
        """估算FPS"""
        model_info = self.MODELS[model]
        platform_info = self.PLATFORMS[platform]
        
        # 有效算力 (考虑量化加速)
        effective_tops = platform_info['tops']
        if quantization == 'int8':
            effective_tops *= 2  # INT8约2倍加速
        elif quantization == 'fp16':
            effective_tops *= 1.0
        
        # 利用率 (边缘设备通常30-60%)
        utilization = 0.4
        effective_flops = effective_tops * 1e12 * utilization
        
        # FPS = 有效算力 / 模型FLOPs
        fps = effective_flops / (model_info['flops'] * 1e9)
        
        # 内存带宽限制
        if 'Nano' in platform and model in ['YOLOv8l', 'RT-DETR-l']:
            fps *= 0.5  # 内存受限
        
        return round(fps, 1)
    
    def estimate_power(self, model: str, platform: str) -> float:
        """估算功耗"""
        platform_info = self.PLATFORMS[platform]
        model_info = self.MODELS[model]
        
        # 功耗与算力利用率成正比
        utilization = min(1.0, model_info['flops'] / 100)
        min_power, max_power = platform_info['tdp_range']
        
        estimated_power = min_power + (max_power - min_power) * utilization
        return round(estimated_power, 1)
    
    def compute_efficiency(self, model: str, platform: str, quant: str = 'int8') -> dict:
        """计算综合能效"""
        fps = self.estimate_fps(model, platform, quant)
        power = self.estimate_power(model, platform)
        mAP = self.MODELS[model]['mAP']
        
        return {
            'model': model,
            'platform': platform,
            'quantization': quant,
            'mAP': mAP,
            'fps': fps,
            'power_w': power,
            'map_per_watt': round(mAP / power, 2),
            'fps_per_watt': round(fps / power, 2),
            'combined_efficiency': round(mAP * fps / power, 2),
            'meets_30fps': fps >= 30,
            'meets_15fps': fps >= 15,
        }


# 测试
if __name__ == "__main__":
    analyzer = EnergyEfficiencyAnalyzer()
    
    print("=== YOLOv8l INT8 能效对比 ===")
    for platform in analyzer.PLATFORMS:
        result = analyzer.compute_efficiency('YOLOv8l', platform, 'int8')
        print(f"\n{platform}:")
        print(f"  FPS: {result['fps']}, 功耗: {result['power_w']}W")
        print(f"  mAP/W: {result['map_per_watt']}, FPS/W: {result['fps_per_watt']}")
        print(f"  30fps达标: {'✅' if result['meets_30fps'] else '❌'}")
    
    print("\n=== RT-DETR-l INT8 能效对比 ===")
    for platform in analyzer.PLATFORMS:
        result = analyzer.compute_efficiency('RT-DETR-l', platform, 'int8')
        print(f"\n{platform}:")
        print(f"  FPS: {result['fps']}, 功耗: {result['power_w']}W")
        print(f"  综合能效: {result['combined_efficiency']}")

class JetsonOrinNXDeployment:
    """
    Jetson Orin NX DMS部署配置
    
    硬件规格:
    - GPU: Ampere (1024 CUDA cores)
    - NPU: 64 Tensor cores
    - 内存: 16GB LPDDR5 102.4GB/s
    - 功耗模式: 15W / 25W (MAXN)
    """
    
    CONFIG = {
        'power_mode': '25W MAXN',  # 最大性能模式
        'gpu_clock': 765,           # MHz
        'memory_clock': 1600,       # MHz
        'jetson_clocks': True,      # 锁定最大频率
    }
    
    # INT8量化配置
    QUANTIZATION_CONFIG = {
        'calibration_method': 'MinMax',
        'calibration_batches': 200,
        'per_channel': True,
        'weight_bits': 8,
        'activation_bits': 8,
    }
    
    # TensorRT优化配置
    TENSORRT_CONFIG = {
        'workspace_size': 4,  # GB
        'precision': 'INT8',
        'max_batch_size': 1,
        'dynamic_shapes': False,  # 固定shape优化
        'layer_fusion': True,
        'kernel_auto_tuning': True,
        'fp16_fallback': True,  # INT8不支持的层用FP16
    }