import torch
import torch.nn as nn
import torch.nn.functional as F

class CognitiveStateEncoder(nn.Module):
    """
    认知状态编码器
    
    将驾驶员分心状态和驾驶意图编码为条件向量
    
    分心状态:
      - attentive (专注)
      - distracted_phone (手机分心)
      - distracted_cognitive (认知分心/走神)
      
    驾驶意图:
      - straight (直行)
      - left_turn (左转)
      - right_turn (右转)
      - lane_change (变道)
    """
    
    def __init__(self, embed_dim: int = 128):
        super().__init__()
        # 分心状态嵌入
        self.distraction_embed = nn.Embedding(5, embed_dim)
        # 驾驶意图嵌入
        self.intention_embed = nn.Embedding(6, embed_dim)
        # 自动驾驶模式
        self.autopilot_embed = nn.Embedding(2, embed_dim)
        # 融合
        self.fusion = nn.Sequential(
            nn.Linear(embed_dim * 3, embed_dim),
            nn.ReLU(),
            nn.Linear(embed_dim, embed_dim)
        )
        
    def forward(self, distraction: torch.Tensor,
                intention: torch.Tensor,
                autopilot: torch.Tensor) -> torch.Tensor:
        """
        Args:
            distraction: (B,) 分心状态ID
            intention: (B,) 驾驶意图ID
            autopilot: (B,) 自动驾驶模式 (0=manual, 1=autopilot)
            
        Returns:
            cognitive_feat: (B, embed_dim) 认知状态特征
        """
        d = self.distraction_embed(distraction)
        i = self.intention_embed(intention)
        a = self.autopilot_embed(autopilot)
        return self.fusion(torch.cat([d, i, a], dim=-1))


class CondConvAttentionNet(nn.Module):
    """
    条件卷积注意力预测网络
    
    基于CoCAtt论文提出的CondConv基线:
    - 边缘场景图 → 视觉特征
    - 认知状态 → 条件卷积核
    - 条件化注意力预测
    
    论文Section 4: Cognitive-Conditioned Attention Learning
    """
    
    def __init__(self, embed_dim: int = 128):
        super().__init__()
        self.cognitive_encoder = CognitiveStateEncoder(embed_dim)
        
        # 场景视觉编码
        self.scene_encoder = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(128, 256, 3, stride=2, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d((8, 16))
        )
        
        # CondConv: 条件卷积核
        self.cond_conv = nn.Sequential(
            nn.Linear(embed_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 256 * 3 * 3)  # 条件化卷积参数
        )
        
        # 注意力解码
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, 
                              output_padding=1), nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1,
                              output_padding=1), nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1,
                              output_padding=1), nn.ReLU(),
            nn.ConvTranspose2d(32, 1, 3, stride=2, padding=1,
                              output_padding=1),
            nn.Sigmoid()
        )
        
    def forward(self, scene: torch.Tensor, distraction: torch.Tensor,
                intention: torch.Tensor, autopilot: torch.Tensor) -> torch.Tensor:
        """
        Args:
            scene: (B, 3, 256, 512) 驾驶场景图像
            distraction: (B,) 分心状态
            intention: (B,) 驾驶意图
            autopilot: (B,) 自动驾驶模式
            
        Returns:
            attention_map: (B, 1, 256, 512) 预测注意力热力图
        """
        # 认知状态编码
        cog_feat = self.cognitive_encoder(distraction, intention, autopilot)
        
        # 场景特征
        scene_feat = self.scene_encoder(scene)  # (B, 256, 8, 16)
        
        # CondConv: 条件化卷积核
        conv_params = self.cond_conv(cog_feat)  # (B, 256*3*3)
        conv_weight = conv_params.view(-1, 256, 3, 3)  # 条件卷积核
        
        # 应用条件卷积（分组卷积模拟）
        # 每个样本用自己的卷积核
        B, C, H, W = scene_feat.shape
        conditioned_feat = torch.zeros_like(scene_feat)
        for i in range(B):
            conditioned_feat[i:i+1] = F.conv2d(
                scene_feat[i:i+1], conv_weight[i:i+1], padding=1
            )
        
        conditioned_feat = F.relu(conditioned_feat)
        
        # 解码为注意力图
        attention_map = self.decoder(conditioned_feat)
        
        # 插值到原始尺寸
        attention_map = F.interpolate(
            attention_map, size=(256, 512), mode='bilinear',
            align_corners=False
        )
        
        return attention_map


class GazeCalibrationNet(nn.Module):
    """
    注视校准网络
    
    CoCAtt提出的coarse-to-fine方法:
    将低分辨率webcam注视数据"对齐"到高分辨率
    
    论文Section 4.2: Real-time Gaze Calibration
    """
    
    def __init__(self):
        super().__init__()
        # 粗注视编码
        self.coarse_encoder = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
        )
        # 场景上下文编码
        self.scene_encoder = nn.Sequential(
            nn.Conv2d(3, 16, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(16, 32, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
        )
        # 融合 + 精细化
        self.refine = nn.Sequential(
            nn.Conv2d(96, 64, 3, padding=1), nn.ReLU(),
            nn.Conv2d(64, 32, 3, padding=1), nn.ReLU(),
            nn.Conv2d(32, 1, 3, padding=1),
            nn.Sigmoid()
        )
        
    def forward(self, coarse_gaze: torch.Tensor,
                scene: torch.Tensor) -> torch.Tensor:
        """
        Args:
            coarse_gaze: (B, 1, 64, 128) 低分辨率注视图
            scene: (B, 3, 256, 512) 场景图像
            
        Returns:
            refined_gaze: (B, 1, 64, 128) 校准后的高分辨率注视图
        """
        gaze_feat = self.coarse_encoder(coarse_gaze)
        scene_feat = self.scene_encoder(scene)
        
        # 确保尺寸匹配
        target_h, target_w = gaze_feat.shape[2:]
        scene_feat = F.interpolate(
            scene_feat, size=(target_h, target_w), mode='bilinear',
            align_corners=False
        )
        
        fused = torch.cat([gaze_feat, scene_feat], dim=1)
        refined = self.refine(fused)
        
        return refined


# ============ 完整测试代码 ============
if __name__ == "__main__":
    import numpy as np
    import time
    
    BATCH_SIZE = 8
    
    # 初始化模型
    model = CondConvAttentionNet(embed_dim=128)
    calibrator = GazeCalibrationNet()
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"CondConvAttentionNet参数量: {total_params:,}")
    
    # 模拟CoCAtt数据
    np.random.seed(42)
    
    # 认知状态定义
    DISTRACTIONS = {0: 'attentive', 1: 'phone', 2: 'cognitive', 3: 'visual', 4: 'mixed'}
    INTENTIONS = {0: 'straight', 1: 'left', 2: 'right', 3: 'lane_change', 4: 'stop', 5: 'park'}
    
    def generate_cocatt_sample():
        """生成一个CoCAtt样本"""
        scene = torch.rand(3, 256, 512)
        distraction = torch.randint(0, 5, (1,)).item()
        intention = torch.randint(0, 6, (1,)).item()
        autopilot = torch.randint(0, 2, (1,)).item()
        
        # 生成注意力图（根据认知状态调整）
        attention = torch.zeros(1, 256, 512)
        if distraction == 0:  # 专注
            if intention == 0:  # 直行
                attention[:, 100:180, 200:350] = 0.8  # 前方
            elif intention == 1:  # 左转
                attention[:, 100:180, 100:250] = 0.7  # 左侧
                attention[:, 50:120, 200:350] = 0.5    # 镜像
        elif distraction == 1:  # 手机
            attention[:, 150:200, 350:500] = 0.9  # 右下方
        elif distraction == 2:  # 认知分心
            attention[:, 100:180, 200:350] = 0.4  # 模糊前方
            attention += torch.rand_like(attention) * 0.1  # 噪声
            
        return scene, distraction, intention, autopilot, attention.squeeze(0)
    
    # 生成batch
    print("\n--- 生成CoCAtt模拟数据 ---")
    scenes, all_dist, all_intent, all_auto, all_gt = [], [], [], [], []
    for _ in range(BATCH_SIZE):
        s, d, i, a, gt = generate_cocatt_sample()
        scenes.append(s)
        all_dist.append(d)
        all_intent.append(i)
        all_auto.append(a)
        all_gt.append(gt)
    
    scene_batch = torch.stack(scenes)
    dist_batch = torch.tensor(all_dist)
    intent_batch = torch.tensor(all_intent)
    auto_batch = torch.tensor(all_auto)
    gt_batch = torch.stack(all_gt)
    
    # 训练
    print("\n--- 训练CondConv注意力模型 ---")
    criterion = nn.KLDivLoss(reduction='batchmean')
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    
    model.train()
    for epoch in range(3):
        optimizer.zero_grad()
        pred = model(scene_batch, dist_batch, intent_batch, auto_batch)
        
        # KL散度
        log_pred = F.log_softmax(pred.view(BATCH_SIZE, -1), dim=1)
        target = F.softmax(gt_batch.view(BATCH_SIZE, -1), dim=1)
        loss = criterion(log_pred, target)
        
        loss.backward()
        optimizer.step()
        
        # CC (Correlation Coefficient)
        pred_flat = pred.view(BATCH_SIZE, -1)
        gt_flat = gt_batch.view(BATCH_SIZE, -1)
        pred_mean = pred_flat.mean(dim=1, keepdim=True)
        gt_mean = gt_flat.mean(dim=1, keepdim=True)
        cc = ((pred_flat - pred_mean) * (gt_flat - gt_mean)).sum(dim=1) / \
             (torch.norm(pred_flat - pred_mean, dim=1) * 
              torch.norm(gt_flat - gt_mean, dim=1) + 1e-8)
        
        print(f"Epoch {epoch+1}: KLD={loss.item():.4f}, CC={cc.mean().item():.4f}")
    
    # 认知状态影响分析
    print("\n--- 认知状态影响分析 ---")
    model.eval()
    with torch.no_grad():
        for dist_id, dist_name in DISTRACTIONS.items():
            d = torch.full((1,), dist_id, dtype=torch.long)
            i = torch.zeros(1, dtype=torch.long)  # 直行
            a = torch.zeros(1, dtype=torch.long)  # 手动
            
            pred = model(scene_batch[:1], d, i, a)
            
            # 计算注视分布
            pred_flat = pred.view(-1)
            # 分区: 前方/左/右/镜像/下方
            front = pred[0, 0, 100:180, 200:350].sum().item()
            left = pred[0, 0, 100:180, :200].sum().item()
            right = pred[0, 0, 100:180, 350:].sum().item()
            mirror = pred[0, 0, :100, :].sum().item()
            bottom = pred[0, 0, 180:, :].sum().item()
            total = front + left + right + mirror + bottom + 1e-8
            
            print(f"  {dist_name:15s}: front={front/total:.1%}, "
                  f"left={left/total:.1%}, right={right/total:.1%}, "
                  f"mirror={mirror/total:.1%}, bottom={bottom/total:.1%}")
    
    # 注视校准测试
    print("\n--- 注视校准测试 ---")
    coarse_gaze = torch.rand(BATCH_SIZE, 1, 64, 128) * 0.3
    refined = calibrator(coarse_gaze, scene_batch)
    print(f"  Coarse: {coarse_gaze.shape}, mean={coarse_gaze.mean():.4f}")
    print(f"  Refined: {refined.shape}, mean={refined.mean():.4f}")
    
    # 推理速度
    with torch.no_grad():
        start = time.time()
        for _ in range(50):
            _ = model(scene_batch[:1], dist_batch[:1], intent_batch[:1], auto_batch[:1])
        elapsed = (time.time() - start) / 50
        print(f"\n单样本推理延迟: {elapsed*1000:.1f}ms")