import torch
import torch.nn as nn
import numpy as np

class CANTransformerDetector(nn.Module):
    """
    CAN总线异常检测Transformer
    
    输入：CAN帧序列 (ID, Data[8], Timestamp)
    输出：异常评分 (0=正常, 1=异常)
    
    核心设计：
    1. CAN ID嵌入 — 学习不同ID的语义
    2. Payload编码 — 8字节数据特征
    3. 时序Transformer — 捕获帧间时序关系
    4. 异常头 — 重构误差+分类双任务
    """
    
    def __init__(self, 
                 n_can_ids: int = 512,
                 d_model: int = 128,
                 n_heads: int = 4,
                 n_layers: int = 3,
                 payload_len: int = 8,
                 max_seq_len: int = 100):
        super().__init__()
        
        # CAN ID嵌入
        self.id_embedding = nn.Embedding(n_can_ids, d_model)
        
        # Payload编码器 (8字节 → d_model)
        self.payload_encoder = nn.Sequential(
            nn.Linear(payload_len * 8, 128),  # 每字节8bit
            nn.GELU(),
            nn.Linear(128, d_model)
        )
        
        # 时间间隔编码
        self.time_encoder = nn.Linear(1, d_model)
        
        # 位置编码
        self.pos_encoding = self._create_pos_encoding(max_seq_len, d_model)
        
        # Transformer编码器
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=d_model * 4,
            dropout=0.1,
            batch_first=True,
            activation='gelu'
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, n_layers)
        
        # 重构头（自编码器模式）
        self.reconstruct_head = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, payload_len * 8)  # 重构payload
        )
        
        # 异常分类头
        self.anomaly_head = nn.Sequential(
            nn.Linear(d_model, d_model // 2),
            nn.GELU(),
            nn.Linear(d_model // 2, 1),
            nn.Sigmoid()
        )
    
    def _create_pos_encoding(self, max_len, d_model):
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).float().unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * 
            (-np.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        return pe.unsqueeze(0)
    
    def forward(self, can_ids, payloads, time_deltas):
        """
        Args:
            can_ids: (B, T) CAN ID序列
            payloads: (B, T, 8) 8字节数据
            time_deltas: (B, T, 1) 帧间时间间隔(ms)
        
        Returns:
            anomaly_scores: (B, T) 每帧异常评分
            reconstructed: (B, T, 64) 重构的payload
        """
        B, T = can_ids.shape
        
        # 嵌入
        id_feat = self.id_embedding(can_ids)  # (B, T, d)
        payload_bits = (payloads.float() > 0.5).float().reshape(B, T, -1)
        payload_feat = self.payload_encoder(payload_bits)  # (B, T, d)
        time_feat = self.time_encoder(time_deltas)  # (B, T, d)
        
        # 融合
        x = id_feat + payload_feat + time_feat
        x = x + self.pos_encoding[:, :T, :].to(x.device)
        
        # Transformer
        x = self.transformer(x)
        
        # 输出
        anomaly_scores = self.anomaly_head(x).squeeze(-1)  # (B, T)
        reconstructed = self.reconstruct_head(x)  # (B, T, 64)
        
        return anomaly_scores, reconstructed


# 测试
if __name__ == "__main__":
    model = CANTransformerDetector()
    
    # 模拟CAN总线数据
    B, T = 2, 50  # 2个序列，各50帧
    can_ids = torch.randint(0, 512, (B, T))
    payloads = torch.randint(0, 256, (B, T, 8)).float()
    time_deltas = torch.randn(B, T, 1) * 10 + 10  # ~10ms间隔
    
    scores, recon = model(can_ids, payloads, time_deltas)
    
    print(f"输入: CAN ID {can_ids.shape}, Payload {payloads.shape}")
    print(f"异常评分: {scores.shape}")
    print(f"重构Payload: {recon.shape}")
    print(f"\n示例评分: {scores[0, :5].detach().numpy()}")

from dataclasses import dataclass
from typing import List

@dataclass
class FederatedConfig:
    """
    联邦学习配置
    
    核心原则：
    - 原始CAN数据永不离开车辆
    - 只传输模型参数（梯度/权重）
    - 服务器聚合参数，分发全局模型
    """
    n_vehicles: int = 100  # 参与车辆数
    local_epochs: int = 5   # 本地训练轮次
    global_rounds: int = 50  # 全局聚合轮次
    aggregation: str = "fedavg"  # 聚合策略
    privacy: str = "dp"     # 差分隐私


class FederatedCANTrainer:
    """联邦学习CAN异常检测训练器"""
    
    def __init__(self, config: FederatedConfig):
        self.config = config
        self.global_model = CANTransformerDetector()
    
    def local_train(self, vehicle_data, local_model, epochs=5):
        """
        车辆端本地训练
        
        原始CAN数据保留在车内
        只输出模型参数更新
        """
        optimizer = torch.optim.Adam(local_model.parameters(), lr=1e-4)
        
        for epoch in range(epochs):
            for batch in vehicle_data:
                can_ids, payloads, time_deltas, labels = batch
                
                # 前向
                scores, recon = local_model(can_ids, payloads, time_deltas)
                
                # 双任务损失
                cls_loss = nn.BCELoss()(scores, labels.float())
                recon_loss = nn.MSELoss()(recon, payloads.float().reshape(*recon.shape))
                loss = cls_loss + 0.1 * recon_loss
                
                # 反向
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
        
        # 返回参数更新（不返回数据）
        return local_model.state_dict()
    
    def aggregate(self, client_updates: List[dict]):
        """
        服务器端FedAvg聚合
        
        简单加权平均各车辆的模型参数
        """
        global_state = self.global_model.state_dict()
        
        for key in global_state:
            # 平均所有客户端的参数
            avg = torch.mean(
                torch.stack([update[key] for update in client_updates]),
                dim=0
            )
            global_state[key] = avg
        
        self.global_model.load_state_dict(global_state)
        return self.global_model.state_dict()
    
    def add_dp_noise(self, params, epsilon=1.0):
        """
        差分隐私：在参数上添加高斯噪声
        
        防止从模型参数反推CAN数据
        """
        for key in params:
            if params[key].is_floating_point():
                noise = torch.randn_like(params[key]) * (1.0 / epsilon)
                params[key] += noise
        return params


# 模拟联邦学习
if __name__ == "__main__":
    config = FederatedConfig(n_vehicles=3, local_epochs=2)
    trainer = FederatedCANTrainer(config)
    
    # 模拟3辆车各自训练
    print(f"参与车辆: {config.n_vehicles}")
    print(f"本地训练轮次: {config.local_epochs}")
    print(f"全局聚合轮次: {config.global_rounds}")
    print(f"聚合策略: {config.aggregation}")
    print(f"隐私保护: {config.privacy}")
    print(f"\n流程: 车辆本地训练 → 上传参数(非数据) → 服务器聚合 → 分发新模型")
    print(f"优势: 原始CAN数据永不离开车辆，保护车主隐私")

class CANAttackScenarios:
    """
    CAN总线攻击场景与检测
    
    IMS相关攻击：
    1. DMS欺骗 — 注入伪造的"正常驾驶"CAN帧
    2. 疲劳注入 — 注入伪造的疲劳信号导致误报
    3. OMS阻断 — 洪泛CPD相关CAN ID
    4. 安全带旁路 — 伪造安全带已系信号
    """
    
    scenarios = {
        "flood_attack": {
            "desc": "洪泛攻击：高频发送同一CAN ID",
            "detection": "Transformer检测到异常高的帧频率",
            "can_id": "0x123",
            "normal_rate": "10ms间隔",
            "attack_rate": "1ms间隔",
        },
        "spoof_attack": {
            "desc": "欺骗攻击：伪造DMS状态帧",
            "detection": "Payload重构误差增大",
            "can_id": "0x456",
            "normal_payload": "DMS_STATUS: NORMAL",
            "attack_payload": "DMS_STATUS: NORMAL (但时序不匹配)",
        },
        "replay_attack": {
            "desc": "重放攻击：重放历史正常CAN帧",
            "detection": "时序模式与当前上下文不匹配",
            "can_id": "0x789",
            "normal": "真实时序",
            "attack": "历史时序（时间戳不连续）",
        },
        "masquerade_attack": {
            "desc": "伪装攻击：冒充其他ECU发送帧",
            "detection": "Transformer检测到ID-Payload组合不匹配",
            "can_id": "0x1AB",
            "normal": "ECU_A发送特定payload模式",
            "attack": "ECU_B发送不同payload模式",
        }
    }

for name, scenario in CANAttackScenarios.scenarios.items():
    print(f"\n{'='*50}")
    print(f"攻击: {name}")
    print(f"描述: {scenario['desc']}")
    print(f"检测: {scenario['detection']}")
    print(f"CAN ID: {scenario['can_id']}")