"""
aiMotive aiWare5 × Cadence NeuroEdge 逐层分配概念

关键设计：
1. 每层可标注执行引擎（aiWare5 或 NeuroEdge）
2. 标注是显式、手动的（可记录、可复现）
3. 不受工具链边界约束
4. 支持随机分配验证鲁棒性
"""

from typing import List, Dict
from dataclasses import dataclass

@dataclass
class NetworkLayer:
    """神经网络层定义"""
    name: str
    layer_type: str  # conv, attention, ffn, bn, fc, etc.
    params: Dict
    target_engine: str = "aiWare5"  # 或 "NeuroEdge"

class HeterogeneousScheduler:
    """
    异构调度器：在aiWare5和NeuroEdge间分配网络层
    
    设计原则：
    1. 显式标注（非自动推断）
    2. 可复现（相同标注→相同执行）
    3. 支持任意分配策略
    """
    
    def __init__(self, engines: List[str]):
        self.engines = engines
        self.layers: List[NetworkLayer] = []
        self.assignment: Dict[str, str] = {}  # layer_name -> engine
        
    def add_layer(self, layer: NetworkLayer):
        """添加网络层"""
        self.layers.append(layer)
        self.assignment[layer.name] = layer.target_engine
    
    def assign(self, layer_name: str, engine: str):
        """显式分配层到引擎"""
        if engine not in self.engines:
            raise ValueError(f"未知引擎: {engine}")
        self.assignment[layer_name] = engine
        for l in self.layers:
            if l.name == layer_name:
                l.target_engine = engine
    
    def get_execution_plan(self) -> Dict[str, List[str]]:
        """获取执行计划"""
        plan = {engine: [] for engine in self.engines}
        for layer in self.layers:
            engine = self.assignment.get(layer.name, "aiWare5")
            plan[engine].append(layer.name)
        return plan
    
    def estimate_performance(self, 
                             aiware5_throughput=100,  # TOPS
                             neuroedge_throughput=30):  # TOPS
        """
        估算性能（简化模型）
        
        Args:
            aiware5_throughput: aiWare5算力 (TOPS)
            neuroedge_throughput: NeuroEdge算力 (TOPS)
        """
        plan = self.get_execution_plan()
        
        # 简化：每层计算量按FLOPs估算
        total_flops = 0
        aiware5_flops = 0
        neuroedge_flops = 0
        
        for layer in self.layers:
            flops = layer.params.get('flops', 1e9)
            total_flops += flops
            engine = self.assignment.get(layer.name, "aiWare5")
            if engine == "aiWare5":
                aiware5_flops += flops
            else:
                neuroedge_flops += flops
        
        aiware5_time = aiware5_flops / (aiware5_throughput * 1e12)
        neuroedge_time = neuroedge_flops / (neuroedge_throughput * 1e12)
        
        # 总时间 = max(两引擎) + 通信开销
        comm_overhead = len(plan["NeuroEdge"]) * 0.001  # 每层切换1ms
        
        return {
            'total_flops': total_flops,
            'aiware5_flops': aiware5_flops,
            'neuroedge_flops': neuroedge_flops,
            'aiware5_time_s': aiware5_time,
            'neuroedge_time_s': neuroedge_time,
            'comm_overhead_s': comm_overhead,
            'estimated_latency_s': max(aiware5_time, neuroedge_time) + comm_overhead
        }


# 测试：模拟DETR分配
if __name__ == "__main__":
    print("=" * 60)
    print("aiWare5 × NeuroEdge 逐层分配模拟")
    print("=" * 60)
    
    scheduler = HeterogeneousScheduler(["aiWare5", "NeuroEdge"])
    
    # 构建简化DETR网络
    layers_def = [
        ("backbone_conv1", "conv", {"flops": 2e9}),
        ("backbone_conv2", "conv", {"flops": 4e9}),
        ("backbone_conv3", "conv", {"flops": 8e9}),
        ("backbone_conv4", "conv", {"flops": 16e9}),
        ("encoder_attn1", "attention", {"flops": 3e9}),
        ("encoder_attn2", "attention", {"flops": 3e9}),
        ("decoder_attn1", "attention", {"flops": 2e9}),
        ("decoder_ffn1", "ffn", {"flops": 1e9}),
        ("head_fc", "fc", {"flops": 0.5e9}),
    ]
    
    for name, ltype, params in layers_def:
        scheduler.add_layer(NetworkLayer(name, ltype, params))
    
    # 策略1: 全部在aiWare5
    print("\n策略1: 全部aiWare5")
    for name, _, _ in layers_def:
        scheduler.assign(name, "aiWare5")
    perf1 = scheduler.estimate_performance()
    print(f"  总FLOPs: {perf1['total_flops']/1e9:.1f}G")
    print(f"  aiWare5时间: {perf1['aiware5_time_s']*1000:.2f}ms")
    print(f"  预估延迟: {perf1['estimated_latency_s']*1000:.2f}ms")
    
    # 策略2: Attention放NeuroEdge, 其余aiWare5
    print("\n策略2: Attention→NeuroEdge, 其余→aiWare5")
    for name, ltype, _ in layers_def:
        if ltype == "attention":
            scheduler.assign(name, "NeuroEdge")
        else:
            scheduler.assign(name, "aiWare5")
    perf2 = scheduler.estimate_performance()
    print(f"  aiWare5 FLOPs: {perf2['aiware5_flops']/1e9:.1f}G → {perf2['aiware5_time_s']*1000:.2f}ms")
    print(f"  NeuroEdge FLOPs: {perf2['neuroedge_flops']/1e9:.1f}G → {perf2['neuroedge_time_s']*1000:.2f}ms")
    print(f"  通信开销: {perf2['comm_overhead_s']*1000:.1f}ms")
    print(f"  预估延迟: {perf2['estimated_latency_s']*1000:.2f}ms")
    
    # 策略3: 交替分配
    print("\n策略3: 交替分配")
    for i, (name, _, _) in enumerate(layers_def):
        engine = "aiWare5" if i % 2 == 0 else "NeuroEdge"
        scheduler.assign(name, engine)
    perf3 = scheduler.estimate_performance()
    plan = scheduler.get_execution_plan()
    print(f"  aiWare5层: {plan['aiWare5']}")
    print(f"  NeuroEdge层: {plan['NeuroEdge']}")
    print(f"  预估延迟: {perf3['estimated_latency_s']*1000:.2f}ms")