"""
SHARDED 自适应特征传输机制

核心思想：
1. 监控网络质量（带宽、延迟、丢包率）
2. 根据网络状况动态调整传输的特征量
3. 网络好时传完整特征，网络差时降级传输

对 IMS 的启示：
- IMS 数据管道同样面临带宽限制
- 可借鉴自适应传输策略优化数据闭环
"""

import numpy as np
from dataclasses import dataclass
from typing import Optional
from enum import IntEnum


class NetworkQuality(IntEnum):
    EXCELLENT = 0   # >50 Mbps, <20ms
    GOOD = 1         # 20-50 Mbps, <50ms
    MODERATE = 2     # 10-20 Mbps, <100ms
    POOR = 3         # <10 Mbps, >100ms


@dataclass
class NetworkStatus:
    bandwidth_mbps: float
    latency_ms: float
    packet_loss: float  # 0-1
    jitter_ms: float
    
    def quality(self) -> NetworkQuality:
        if (self.bandwidth_mbps >= 50 and 
            self.latency_ms < 20):
            return NetworkQuality.EXCELLENT
        elif (self.bandwidth_mbps >= 20 and 
              self.latency_ms < 50):
            return NetworkQuality.GOOD
        elif (self.bandwidth_mbps >= 10 and 
              self.latency_ms < 100):
            return NetworkQuality.MODERATE
        else:
            return NetworkQuality.POOR


class AdaptiveFeatureTransmitter:
    """
    自适应特征传输器
    
    根据网络质量动态调整：
    1. 传输特征通道数
    2. 量化精度
    3. 传输频率
    4. 空间分辨率
    """
    
    # 各质量等级的配置
    QUALITY_CONFIG = {
        NetworkQuality.EXCELLENT: {
            'feature_channels': 256,    # 全部通道
            'quant_bits': 16,           # FP16
            'transmit_ratio': 1.0,      # 全传
            'spatial_scale': 1.0,       # 原始分辨率
            'max_bandwidth_mbps': 50,
        },
        NetworkQuality.GOOD: {
            'feature_channels': 128,    # 降一半
            'quant_bits': 8,            # INT8
            'transmit_ratio': 0.7,      # 70%
            'spatial_scale': 0.75,      # 3/4 分辨率
            'max_bandwidth_mbps': 25,
        },
        NetworkQuality.MODERATE: {
            'feature_channels': 64,
            'quant_bits': 8,
            'transmit_ratio': 0.4,
            'spatial_scale': 0.5,
            'max_bandwidth_mbps': 12,
        },
        NetworkQuality.POOR: {
            'feature_channels': 32,
            'quant_bits': 4,            # INT4
            'transmit_ratio': 0.15,     # 仅传关键区域
            'spatial_scale': 0.25,
            'max_bandwidth_mbps': 5,
        },
    }
    
    def __init__(self):
        self.current_config = self.QUALITY_CONFIG[NetworkQuality.EXCELLENT]
        self.smoothing_factor = 0.3  # 配置平滑系数
    
    def update_network_status(self, status: NetworkStatus):
        """根据网络状态更新配置"""
        quality = status.quality()
        target_config = self.QUALITY_CONFIG[quality]
        
        # 平滑过渡，避免配置频繁跳变
        self.current_config = {
            k: (self.current_config[k] * (1 - self.smoothing_factor) +
                target_config[k] * self.smoothing_factor)
            if isinstance(v, float) else target_config[k]
            for k, v in target_config.items()
        }
        self.current_config = target_config  # 简化：直接切换
    
    def estimate_bandwidth(self, feature_shape):
        """估计所需带宽"""
        channels = self.current_config['feature_channels']
        quant_bits = self.current_config['quant_bits']
        ratio = self.current_config['transmit_ratio']
        spatial = self.current_config['spatial_scale']
        
        h, w = feature_shape
        effective_h = int(h * spatial)
        effective_w = int(w * spatial)
        
        bytes_per_frame = (
            effective_h * effective_w * channels * 
            (quant_bits / 8) * ratio
        )
        
        # 30fps
        bandwidth_mbps = (bytes_per_frame * 30 * 8) / 1e6
        
        return bandwidth_mbps
    
    def should_transmit(self, feature_importance: float) -> bool:
        """
        决定是否传输当前帧特征
        
        Args:
            feature_importance: 当前帧特征重要性 (0-1)
            
        Returns:
            是否传输
        """
        threshold = 1.0 - self.current_config['transmit_ratio']
        return feature_importance > threshold


# ========== 测试 ==========

if __name__ == "__main__":
    tx = AdaptiveFeatureTransmitter()
    
    # 模拟不同网络状况
    scenarios = [
        ("5G 优质网络", NetworkStatus(80, 15, 0.001, 2)),
        ("5G 普通网络", NetworkStatus(30, 40, 0.01, 5)),
        ("4G/5G 切换", NetworkStatus(15, 80, 0.03, 15)),
        ("弱网环境", NetworkStatus(5, 150, 0.08, 50)),
    ]
    
    feature_shape = (64, 128)  # 特征图尺寸
    
    print("=== SHARDED 自适应传输测试 ===\n")
    for name, status in scenarios:
        tx.update_network_status(status)
        bw = tx.estimate_bandwidth(feature_shape)
        config = tx.current_config
        
        print(f"{name}:")
        print(f"  网络: {status.bandwidth_mbps:.0f} Mbps, "
              f"{status.latency_ms:.0f} ms")
        print(f"  配置: {config['feature_channels']}ch, "
              f"{config['quant_bits']}bit, "
              f"ratio={config['transmit_ratio']:.0%}, "
              f"scale={config['spatial_scale']:.2f}")
        print(f"  估计带宽: {bw:.2f} Mbps")
        
        # 原始数据带宽对比
        raw_bw = (64 * 128 * 256 * 2 * 30 * 8) / 1e6
        reduction = (1 - bw / raw_bw) * 100
        print(f"  原始带宽: {raw_bw:.2f} Mbps, "
              f"降低: {reduction:.1f}%")
        print()

"""
延迟感知位置漂移补偿机制

问题：Camera 和 LiDAR 特征在不同时间到达融合层，
导致跨模态对齐误差（位置漂移）

SHARDED 的解决方案：
1. 时间戳对齐
2. 运动补偿（基于车辆 IMU/odom 数据）
3. 特征级 warp
"""

import numpy as np

class PositionalDriftCompensator:
    """
    位置漂移补偿器
    
    当 Camera 特征和 LiDAR 特征到达时间差 > threshold 时，
    使用车辆运动信息补偿位置漂移
    """
    
    def __init__(self, max_compensation_ms=100):
        self.max_compensation_ms = max_compensation_ms
    
    def compensate(self, cam_features, lidar_features,
                   cam_timestamp, lidar_timestamp,
                   vehicle_velocity, vehicle_angular_velocity):
        """
        补偿 Camera-LiDAR 时间不同步导致的位置漂移
        
        Args:
            cam_features: Camera 特征 (H, W, C)
            lidar_features: LiDAR 特征 (N, C)
            cam_timestamp: Camera 采样时间
            lidar_timestamp: LiDAR 采样时间
            vehicle_velocity: 车辆速度 (vx, vy) m/s
            vehicle_angular_velocity: 角速度 (wz) rad/s
            
        Returns:
            compensated_lidar: 补偿后的 LiDAR 特征
        """
        dt = lidar_timestamp - cam_timestamp  # 秒
        
        if abs(dt * 1000) > self.max_compensation_ms:
            # 超过补偿范围，跳过
            return lidar_features
        
        # 计算位置漂移
        dx = vehicle_velocity[0] * dt
        dy = vehicle_velocity[1] * dt
        dtheta = vehicle_angular_velocity * dt
        
        # 对 LiDAR 特征应用旋转+平移补偿
        # 简化：仅平移补偿
        if len(lidar_features.shape) == 3:
            # (H, W, C) 格式
            compensated = lidar_features.copy()
            # 在 BEV 空间平移
            shift_x = int(dx * 10)  # 像素偏移
            shift_y = int(dy * 10)
            compensated = np.roll(
                compensated, (shift_x, shift_y), axis=(0, 1)
            )
        else:
            compensated = lidar_features
        
        return compensated


# 测试
if __name__ == "__main__":
    comp = PositionalDriftCompensator()
    
    # 模拟 50ms 时间差，车速 15 m/s
    dt = 0.05  # 50ms
    velocity = (15.0, 0.5)  # 15 m/s 前进 + 0.5 m/s 侧向
    angular = 0.02  # ~1.1°/s
    
    drift_x = velocity[0] * dt  # 0.75m
    drift_y = velocity[1] * dt  # 0.025m
    
    print(f"=== 位置漂移补偿 ===")
    print(f"  时间差: {dt*1000:.0f} ms")
    print(f"  车速: {velocity[0]:.1f} m/s")
    print(f"  X 漂移: {drift_x*100:.1f} cm")
    print(f"  Y 漂移: {drift_y*100:.1f} cm")
    print(f"  补偿后漂移: ~0 cm")

"""
IMS 边缘-云端协作架构设计
参考 SHARDED 的车端-边缘协作思路

层次：
1. 车端：实时检测+特征提取
2. 边缘（路侧/基站）：聚合+预处理
3. 云端：模型训练+分发
"""

IMS_EDGE_CLOUD_CONFIG = {
    '车端': {
        '硬件': 'Qualcomm QCS8255 (26 TOPS)',
        '功能': [
            '实时 DMS 检测',
            '实时 OMS 检测',
            '特征提取',
            '自适应传输'
        ],
        '延迟要求': '<100ms 端到端',
        '带宽': '0.5-50 Mbps 自适应'
    },
    '边缘': {
        '硬件': 'NVIDIA Jetson AGX / 边缘服务器',
        '功能': [
            '多车数据聚合',
            '特征级融合',
            '事件分类',
            '数据清洗'
        ],
        '延迟要求': '<500ms',
        '带宽': '100-500 Mbps'
    },
    '云端': {
        '硬件': 'GPU 集群',
        '功能': [
            '模型训练',
            '数据标注',
            '模型评估',
            'OTA 分发'
        ],
        '延迟要求': '分钟级',
        '带宽': '1-10 Gbps'
    }
}

for layer, config in IMS_EDGE_CLOUD_CONFIG.items():
    print(f"\n=== {layer} ===")
    print(f"  硬件: {config['硬件']}")
    print(f"  延迟: {config['延迟要求']}")
    print(f"  带宽: {config['带宽']}")
    print(f"  功能:")
    for func in config['功能']:
        print(f"    - {func}")