"""
aiMotive aiWare5 × Cadence NeuroEdge 逐层工作负载分配模拟

演示如何将DMS/OMS神经网络分配到异构NPU+协处理器
"""

from typing import List, Dict
from dataclasses import dataclass

@dataclass
class LayerInfo:
    """网络层信息"""
    name: str
    layer_type: str  # conv, bn, relu, attention, fc
    flops: float     # GFLOPs
    params: float    # M parameters
    input_shape: tuple
    output_shape: tuple

@dataclass
class PlacementResult:
    """分配结果"""
    layer_name: str
    engine: str  # 'aiWare5' or 'NeuroEdge' or 'CPU'
    latency_ms: float
    power_mw: float


class HeterogeneousWorkloadPlacer:
    """
    逐层工作负载分配器
    
    模拟aiMotive的逐层标注分配策略
    """
    
    def __init__(self):
        # aiWare5: 擅长大批量CONV和FC
        self.aiware5_perf = {
            'conv': 4.0,     # TOPS
            'bn': 8.0,       # TOPS
            'relu': 10.0,   # TOPS
            'attention': 1.5, # TOPS (较弱)
            'fc': 6.0,       # TOPS
        }
        self.aiware5_power = {
            'conv': 2.5, 'bn': 1.0, 'relu': 0.5,
            'attention': 3.0, 'fc': 2.0
        }
        
        # NeuroEdge: 擅长attention和轻量计算
        self.neuroedge_perf = {
            'conv': 1.5,
            'bn': 4.0,
            'relu': 5.0,
            'attention': 6.0,  # TOPS (较强)
            'fc': 3.0,
        }
        self.neuroedge_power = {
            'conv': 2.0, 'bn': 0.8, 'relu': 0.3,
            'attention': 1.5, 'fc': 1.5
        }
    
    def estimate_latency(self, layer: LayerInfo, engine: str) -> float:
        """估算单层延迟"""
        if engine == 'aiWare5':
            perf = self.aiware5_perf.get(layer.layer_type, 2.0)
            power = self.aiware5_power.get(layer.layer_type, 2.0)
        elif engine == 'NeuroEdge':
            perf = self.neuroedge_perf.get(layer.layer_type, 2.0)
            power = self.neuroedge_power.get(layer.layer_type, 2.0)
        else:
            perf = 0.5  # CPU
            power = 5.0
        
        latency = layer.flops / perf * 1000  # ms
        return latency
    
    def estimate_power(self, layer: LayerInfo, engine: str) -> float:
        """估算单层功耗"""
        if engine == 'aiWare5':
            return self.aiware5_power.get(layer.layer_type, 2.0) * layer.flops
        elif engine == 'NeuroEdge':
            return self.neuroedge_power.get(layer.layer_type, 2.0) * layer.flops
        return 5.0 * layer.flops
    
    def place_greedy(self, layers: List[LayerInfo]) -> List[PlacementResult]:
        """贪心分配：每层选择延迟更低的引擎"""
        results = []
        for layer in layers:
            lat_aim = self.estimate_latency(layer, 'aiWare5')
            lat_neuro = self.estimate_latency(layer, 'NeuroEdge')
            
            if lat_aim <= lat_neuro:
                engine = 'aiWare5'
                latency = lat_aim
            else:
                engine = 'NeuroEdge'
                latency = lat_neuro
            
            power = self.estimate_power(layer, engine)
            results.append(PlacementResult(
                layer_name=layer.name,
                engine=engine,
                latency_ms=latency,
                power_mw=power
            ))
        return results
    
    def place_balanced(self, layers: List[LayerInfo]) -> List[PlacementResult]:
        """平衡分配：考虑引擎负载均衡"""
        results = []
        aim_load = 0.0
        neuro_load = 0.0
        
        for layer in layers:
            lat_aim = self.estimate_latency(layer, 'aiWare5')
            lat_neuro = self.estimate_latency(layer, 'NeuroEdge')
            
            # 考虑当前负载
            adjusted_aim = lat_aim + aim_load * 0.1
            adjusted_neuro = lat_neuro + neuro_load * 0.1
            
            if adjusted_aim <= adjusted_neuro:
                engine = 'aiWare5'
                aim_load += lat_aim
                latency = lat_aim
            else:
                engine = 'NeuroEdge'
                neuro_load += lat_neuro
                latency = lat_neuro
            
            power = self.estimate_power(layer, engine)
            results.append(PlacementResult(
                layer_name=layer.name,
                engine=engine,
                latency_ms=latency,
                power_mw=power
            ))
        return results


# DMS模型模拟
def create_dms_model_layers() -> List[LayerInfo]:
    """创建典型DMS模型的层列表"""
    return [
        LayerInfo('conv1', 'conv', 0.3, 0.01, (3,224,224), (64,112,112)),
        LayerInfo('bn1', 'bn', 0.05, 0.001, (64,112,112), (64,112,112)),
        LayerInfo('relu1', 'relu', 0.02, 0, (64,112,112), (64,112,112)),
        LayerInfo('conv2', 'conv', 1.2, 0.05, (64,112,112), (128,56,56)),
        LayerInfo('bn2', 'bn', 0.1, 0.002, (128,56,56), (128,56,56)),
        LayerInfo('relu2', 'relu', 0.04, 0, (128,56,56), (128,56,56)),
        LayerInfo('attn1', 'attention', 2.5, 0.1, (128,56,56), (128,56,56)),
        LayerInfo('conv3', 'conv', 0.8, 0.03, (128,56,56), (256,28,28)),
        LayerInfo('bn3', 'bn', 0.15, 0.003, (256,28,28), (256,28,28)),
        LayerInfo('attn2', 'attention', 3.0, 0.15, (256,28,28), (256,28,28)),
        LayerInfo('fc1', 'fc', 0.5, 0.02, (256,7,7), (512,)),
        LayerInfo('fc2', 'fc', 0.1, 0.001, (512,), (2,)),  # 二分类: 疲劳/正常
    ]

# 测试
if __name__ == "__main__":
    print("=" * 70)
    print("aiMotive aiWare5 × Cadence NeuroEdge 工作负载分配模拟")
    print("=" * 70)
    
    layers = create_dms_model_layers()
    placer = HeterogeneousWorkloadPlacer()
    
    # 贪心分配
    print("\n--- 贪心分配（每层选最快引擎）---")
    results_greedy = placer.place_greedy(layers)
    total_latency_g = sum(r.latency_ms for r in results_greedy)
    total_power_g = sum(r.power_mw for r in results_greedy)
    aim_count_g = sum(1 for r in results_greedy if r.engine == 'aiWare5')
    neuro_count_g = sum(1 for r in results_greedy if r.engine == 'NeuroEdge')
    
    print(f"{'层名':>10} | {'类型':>10} | {'引擎':>12} | {'延迟(ms)':>10} | {'功耗(mW)':>10}")
    print("-" * 62)
    for layer, result in zip(layers, results_greedy):
        print(f"{result.layer_name:>10} | {layer.layer_type:>10} | {result.engine:>12} | {result.latency_ms:10.3f} | {result.power_mw:10.3f}")
    
    print(f"\n总延迟: {total_latency_g:.3f} ms")
    print(f"总功耗: {total_power_g:.3f} mW")
    print(f"aiWare5层: {aim_count_g}, NeuroEdge层: {neuro_count_g}")
    
    # 平衡分配
    print(f"\n--- 平衡分配（考虑负载均衡）---")
    results_balanced = placer.place_balanced(layers)
    total_latency_b = sum(r.latency_ms for r in results_balanced)
    total_power_b = sum(r.power_mw for r in results_balanced)
    aim_count_b = sum(1 for r in results_balanced if r.engine == 'aiWare5')
    neuro_count_b = sum(1 for r in results_balanced if r.engine == 'NeuroEdge')
    
    print(f"{'层名':>10} | {'引擎':>12} | {'延迟(ms)':>10} | {'功耗(mW)':>10}")
    print("-" * 50)
    for layer, result in zip(layers, results_balanced):
        print(f"{result.layer_name:>10} | {result.engine:>12} | {result.latency_ms:10.3f} | {result.power_mw:10.3f}")
    
    print(f"\n总延迟: {total_latency_b:.3f} ms")
    print(f"总功耗: {total_power_b:.3f} mW")
    print(f"aiWare5层: {aim_count_b}, NeuroEdge层: {neuro_count_b}")
    
    # 对比
    print(f"\n--- 策略对比 ---")
    print(f"{'策略':>12} | {'总延迟(ms)':>12} | {'总功耗(mW)':>12} | {'均衡度':>8}")
    print("-" * 52)
    print(f"{'贪心':>12} | {total_latency_g:12.3f} | {total_power_g:12.3f} | {aim_count_g}:{neuro_count_g}")
    print(f"{'平衡':>12} | {total_latency_b:12.3f} | {total_power_b:12.3f} | {aim_count_b}:{neuro_count_b}")

# 逐层标注示例（伪代码，参考aiMotive工具链）
# 每层可以显式标注执行引擎

model = load_model("dms_model.onnx")

# 逐层标注
for layer in model.layers:
    if layer.type in ['Conv', 'BatchNorm', 'ReLU']:
        layer.target = 'aiWare5'
    elif layer.type in ['Attention', 'GEMM']:
        layer.target = 'NeuroEdge'
    else:
        layer.target = 'auto'  # 工具链自动选择

# 编译（两个工具链各自编译标注的层）
aiware5_compile(model.layers_for('aiWare5'))
neuroedge_compile(model.layers_for('NeuroEdge'))

# 运行时调度（自动，无需自定义调度器）
runtime = HeterogeneousRuntime(aiware5_engine, neuroedge_engine)
output = runtime.run(model, input_frame)