"""
基于深度图像的3D人体姿态估计

论文: In-vehicle 3D vision for perceiving dangerous driving behaviors
期刊: Scientific Reports (Nature), 2026

核心: 从 ToF 深度图像直接回归 16 个 3D 关节点
关键创新: 模型比基准小 70%，准确率保持 96.02%
"""

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

class DepthPoseNet(nn.Module):
    """
    轻量化 3D 姿态估计网络
    
    输入: 深度图像 (B, 1, H, W)  
    输出: 16 个 3D 关节点坐标 (B, 16, 3)
    
    相比基准模型（如 PoseFormer/SimpleBaseline）体积减小 70%
    """
    
    # 16 个关键关节点定义
    JOINT_NAMES = [
        'head_top', 'neck',          # 头部
        'right_shoulder', 'right_elbow', 'right_wrist',  # 右臂
        'left_shoulder', 'left_elbow', 'left_wrist',     # 左臂
        'right_hip', 'right_knee', 'right_ankle',        # 右腿
        'left_hip', 'left_knee', 'left_ankle',           # 左腿
        'pelvis', 'spine'                                # 躯干
    ]
    
    def __init__(self, num_joints=16, depth_channels=1):
        super().__init__()
        
        # 轻量特征提取骨干
        self.backbone = nn.Sequential(
            # Stage 1: 224x224 → 56x56
            self._conv_block(depth_channels, 32, 3, 2, 1),
            self._conv_block(32, 64, 3, 2, 1),
            # Stage 2: 56x56 → 28x28
            self._conv_block(64, 96, 3, 2, 1),
            # Stage 3: 28x28 → 14x14
            self._conv_block(96, 128, 3, 2, 1),
            # Stage 4: 14x14 → 7x7
            self._conv_block(128, 256, 3, 2, 1),
        )
        
        # 2D 热图预测头 (中间监督导)
        self.heatmap_head = nn.Conv2d(256, num_joints, 1)
        
        # 3D 坐标回归头
        self.regression_head = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(256, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(512, num_joints * 3)
        )
        
        # 关节关系建模（轻量注意力）
        self.joint_attention = JointRelationModule(
            num_joints=num_joints,
            embed_dim=128,
            num_heads=4
        )
    
    def _conv_block(self, in_ch, out_ch, kernel, stride, padding):
        return nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel, stride, padding, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )
    
    def forward(self, depth_image):
        """
        Args:
            depth_image: (B, 1, 224, 224) 归一化深度图
            
        Returns:
            joints_3d: (B, 16, 3) 3D 关节坐标（相机坐标系，毫米）
            heatmaps: (B, 16, 56, 56) 2D 热图（辅助监督）
        """
        # 特征提取
        features = self.backbone(depth_image)  # (B, 256, 7, 7)
        
        # 2D 热图（中间监督）
        heatmaps = self.heatmap_head(features)  # (B, 16, 7, 7)
        heatmaps = F.interpolate(
            heatmaps, scale_factor=8, mode='bilinear',
            align_corners=False
        )  # (B, 16, 56, 56)
        
        # 3D 回归
        joints_3d = self.regression_head(features)  # (B, 16*3)
        joints_3d = joints_3d.view(-1, 16, 3)  # (B, 16, 3)
        
        # 关节关系优化
        joints_3d = self.joint_attention(joints_3d)
        
        return joints_3d, heatmaps


class JointRelationModule(nn.Module):
    """
    关节关系注意力模块
    
    建模关节间的运动学约束关系
    例如: 肘部应在肩部和手腕之间
    """
    def __init__(self, num_joints=16, embed_dim=128, num_heads=4):
        super().__init__()
        self.proj = nn.Linear(3, embed_dim)
        self.attention = nn.MultiheadAttention(
            embed_dim=embed_dim,
            num_heads=num_heads,
            batch_first=True
        )
        self.norm = nn.LayerNorm(embed_dim)
        self.ffn = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 2),
            nn.GELU(),
            nn.Linear(embed_dim * 2, embed_dim),
            nn.Dropout(0.1)
        )
        self.norm2 = nn.LayerNorm(embed_dim)
        self.out_proj = nn.Linear(embed_dim, 3)
        
    def forward(self, joints_3d):
        """
        Args:
            joints_3d: (B, 16, 3) 初始 3D 关节坐标
        Returns:
            refined_joints: (B, 16, 3) 优化后的 3D 关节坐标
        """
        x = self.proj(joints_3d)  # (B, 16, D)
        
        # 自注意力: 建模关节间关系
        attn_out, _ = self.attention(x, x, x)
        x = x + attn_out
        x = self.norm(x)
        
        # FFN
        ffn_out = self.ffn(x)
        x = x + ffn_out
        x = self.norm2(x)
        
        # 回归到 3D 坐标
        refined = self.out_proj(x)  # (B, 16, 3)
        
        # 残差连接
        return joints_3d + refined  # 残差优化


# ============================================
# 损失函数
# ============================================

class PoseLoss(nn.Module):
    """
    组合损失: 3D 坐标 MPJPE + 2D 热图 MSE + 关节关系约束
    """
    def __init__(self, lambda_3d=1.0, lambda_2d=0.5, lambda_bone=0.3):
        super().__init__()
        self.lambda_3d = lambda_3d
        self.lambda_2d = lambda_2d
        self.lambda_bone = lambda_bone
        
    def forward(self, pred_3d, gt_3d, pred_2d, gt_2d):
        """
        Args:
            pred_3d: (B, 16, 3) 预测 3D 关节
            gt_3d: (B, 16, 3) 真值 3D 关节
            pred_2d: (B, 16, 56, 56) 预测 2D 热图
            gt_2d: (B, 16, 56, 56) 真值 2D 热图
        """
        # MPJPE: 平均关节位置误差
        mpjpe = torch.mean(
            torch.norm(pred_3d - gt_3d, dim=-1)
        )  # 标量
        
        # 2D 热图 MSE
        loss_2d = F.mse_loss(pred_2d, gt_2d)
        
        # 骨长一致性约束
        # 定义骨骼连接
        bones = [
            (0, 1),    # head_top → neck
            (1, 2),    # neck → right_shoulder
            (1, 5),    # neck → left_shoulder
            (2, 3),    # right_shoulder → right_elbow
            (3, 4),    # right_elbow → right_wrist
            (5, 6),    # left_shoulder → left_elbow
            (6, 7),    # left_elbow → left_wrist
            (1, 14),   # neck → spine
            (14, 15),  # spine → pelvis
            (15, 8),   # pelvis → right_hip
            (15, 11),  # pelvis → left_hip
        ]
        
        bone_loss = 0
        for j1, j2 in bones:
            pred_len = torch.norm(
                pred_3d[:, j1] - pred_3d[:, j2], dim=-1
            )
            gt_len = torch.norm(
                gt_3d[:, j1] - gt_3d[:, j2], dim=-1
            )
            bone_loss += torch.mean((pred_len - gt_len) ** 2)
        bone_loss /= len(bones)
        
        total = (
            self.lambda_3d * mpjpe +
            self.lambda_2d * loss_2d +
            self.lambda_bone * bone_loss
        )
        
        return total, {
            'mpjpe': mpjpe.item(),
            'loss_2d': loss_2d.item(),
            'bone_loss': bone_loss.item()
        }

"""
ST-GCN++ 驾驶行为识别

基于 3D 骨架序列的时空图卷积网络
输入: (B, T, 16, 3) 3D 关节序列
输出: (B, 10) 10类行为分类
"""

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

class STGCNPlusPlus(nn.Module):
    """
    Spatial-Temporal Graph Convolutional Network++
    
    改进:
    1. 自适应图邻接矩阵（可学习）
    2. 多尺度时间卷积
    3. 注意力增强的图卷积
    """
    
    # 驾驶姿态关节图邻接关系
    ADJACENCY = [
        [0,1], [1,2], [2,3], [3,4],   # 头→右臂
        [1,5], [5,6], [6,7],          # 左臂
        [1,14], [14,15],              # 躯干
        [15,8], [8,9], [9,10],        # 右腿
        [15,11], [11,12], [12,13],    # 左腿
        # 跨肢体连接（提升姿态稳定性）
        [2,5],   # 双肩
        [8,11],  # 双髋
    ]
    
    def __init__(self, num_joints=16, in_channels=3, 
                 num_classes=10, hidden=64):
        super().__init__()
        
        # 数据归一化
        self.input_bn = nn.BatchNorm2d(num_joints * in_channels)
        
        # 多尺度时间卷积
        self.temporal_conv1 = nn.Conv1d(
            num_joints * in_channels, hidden, 
            kernel_size=3, padding=1
        )
        self.temporal_conv2 = nn.Conv1d(
            num_joints * in_channels, hidden,
            kernel_size=5, padding=2
        )
        self.temporal_conv3 = nn.Conv1d(
            num_joints * in_channels, hidden,
            kernel_size=1
        )
        
        # 自适应图卷积层
        self.gcn_layers = nn.ModuleList([
            AdaptiveGCN(hidden, hidden) for _ in range(3)
        ])
        
        # 时间下采样
        self.t_pool = nn.MaxPool1d(kernel_size=2, stride=2)
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool1d(1),
            nn.Flatten(),
            nn.Linear(hidden, hidden),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(hidden, num_classes)
        )
        
    def forward(self, x):
        """
        Args:
            x: (B, T, J, C) 关节序列
               T=时间帧, J=16关节, C=3坐标
        Returns:
            logits: (B, 10) 行为分类
        """
        B, T, J, C = x.shape
        
        # 展平关节和坐标维度
        x = x.reshape(B, T, J * C)  # (B, T, J*C)
        x = x.permute(0, 2, 1)  # (B, J*C, T)
        
        # 归一化
        x = self.input_bn(x.unsqueeze(-1)).squeeze(-1)
        
        # 多尺度时间特征
        t1 = self.temporal_conv1(x)  # 3-frame
        t2 = self.temporal_conv2(x)  # 5-frame
        t3 = self.temporal_conv3(x)  # 1-frame
        x = t1 + t2 + t3  # 融合
        x = F.relu(x)
        
        # 图卷积层
        for gcn in self.gcn_layers:
            x = gcn(x)  # (B, hidden, T)
        
        # 时间下采样
        x = self.t_pool(x)
        
        # 分类
        logits = self.classifier(x)  # (B, 10)
        return logits


class AdaptiveGCN(nn.Module):
    """
    自适应图卷积
    
    可学习邻接矩阵，无需预定义图结构
    """
    def __init__(self, in_ch, out_ch):
        super().__init__()
        # 可学习邻接矩阵
        self.A = nn.Parameter(
            torch.randn(16, 16) * 0.01 + torch.eye(16)
        )
        self.conv = nn.Conv1d(in_ch, out_ch, kernel_size=1)
        self.bn = nn.BatchNorm1d(out_ch)
        
    def forward(self, x):
        """
        Args:
            x: (B, C, T) 时序特征
        Returns:
            out: (B, C, T)
        """
        # 归一化邻接矩阵
        A = F.softmax(self.A, dim=-1)
        
        # 图卷积: 先做空间聚合
        B, C, T = x.shape
        x_reshaped = x.permute(0, 2, 1)  # (B, T, C=J*C_in)
        # 分离关节维度
        # 实际实现中需要 reshape 到 (B, T, J, C_in)
        # 简化: 直接做 1D 卷积
        x = self.conv(x)
        x = self.bn(x)
        x = F.relu(x)
        return x

"""
双视角融合决策模块

解决单视角遮挡问题:
- 前视角: 捕获面部、上肢
- 侧视角: 捕获躯干、手部位置
"""

import torch
import torch.nn as nn

class DualViewFusion(nn.Module):
    """
    双视角检测融合 + 目标检测辅助决策
    """
    def __init__(self, num_classes=10, feature_dim=64):
        super().__init__()
        # 视角置信度学习
        self.view_weight = nn.Sequential(
            nn.Linear(feature_dim * 2, feature_dim),
            nn.ReLU(),
            nn.Linear(feature_dim, 2),
            nn.Softmax(dim=-1)
        )
        # 融合分类
        self.fusion_cls = nn.Sequential(
            nn.Linear(feature_dim * 2, feature_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(feature_dim, num_classes)
        )
        
    def forward(self, feat_front, feat_side, 
                det_front=None, det_side=None):
        """
        Args:
            feat_front: (B, D) 前视角度特征
            feat_side: (B, D) 侧视角度特征
            det_front: 前视目标检测结果（可选）
            det_side: 侧视目标检测结果（可选）
        Returns:
            logits: (B, 10) 融合分类
            weights: (B, 2) 视角权重
        """
        # 拼接特征
        combined = torch.cat([feat_front, feat_side], dim=-1)
        
        # 学习视角权重
        weights = self.view_weight(combined)  # (B, 2)
        
        # 加权融合
        weighted_front = feat_front * weights[:, 0:1]
        weighted_side = feat_side * weights[:, 1:2]
        fused = torch.cat([weighted_front, weighted_side], dim=-1)
        
        # 融合分类
        logits = self.fusion_cls(fused)
        
        # 目标检测辅助决策
        if det_front is not None:
            # 如果检测到手机物体，提升"玩手机"类别置信度
            if 'phone' in det_front:
                logits[:, 2] += 0.5  # 玩手机类boost
            if 'cigarette' in det_front:
                logits[:, 3] += 0.5  # 抽烟类boost
        
        return logits, weights