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

class DualViewToFSystem:
    """
    双视角ToF深度相机系统
    
    主视角：仪表台上方，正对驾驶员
    侧视角：A柱内侧，侧对驾驶员
    
    解决单视角遮挡问题：
    - 手被方向盘遮挡 → 侧视角可见
    - 身体前倾被遮挡 → 侧视角补充
    """
    
    def __init__(self):
        # 主从双视角（论文使用 Azure Kinect ToF）
        self.main_cam = {
            'position': 'dashboard_top_center',
            'resolution': (640, 480),
            'fps': 30,
            'fov': (120, 120),  # H, V degrees
            'sync_error': '<10ms',  # 主从同步误差
        }

        # 侧视角参数
        self.side_cam = {
            'position': 'a_pillar_left',
            'resolution': (640, 480),
            'fps': 30,
            'fov': (120, 120),
            'sync_error': '<10ms',
        }
        
    def get_coverage(self):
        """双视角覆盖区域分析"""
        return {
            'main': ['面部', '上身正面', '双手前方', '大腿正面'],
            'side': ['手臂侧面', '上身侧面', '后背上部', '肩部'],
            'overlap': ['头部', '肩部', '胸部'],  # 双视角均可见
            'blind_spot': ['腰部以下后方', '脚部']
        }


class Lightweight3DPoseEstimator(nn.Module):
    """
    轻量级3D姿态估计器

    论文方案：模型大小约为基准的27%（46.25MB），精度96.02%（基准95.65%）
    输入：ToF深度图像 (1, H, W)
    输出：16个3D关键关节点
    """
    
    # 16个关键关节点定义
    JOINT_NAMES = [
        'head_top', 'head', 'neck',           # 头部
        'left_shoulder', 'right_shoulder',    # 肩部
        'left_elbow', 'right_elbow',          # 肘部
        'left_wrist', 'right_wrist',          # 手腕
        'spine', 'pelvis',                    # 躯干
        'left_hip', 'right_hip',              # 髋部
        'left_knee', 'right_knee',            # 膝盖
    ]
    
    def __init__(self, num_joints: int = 16):
        super().__init__()
        
        # 深度图像特征提取（比RGB更简单，单通道）
        self.backbone = nn.Sequential(
            # Block 1
            nn.Conv2d(1, 32, 3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            # Block 2
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            # Block 3
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            # Block 4
            nn.Conv2d(128, 256, 3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            # Block 5
            nn.Conv2d(256, 512, 3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(),
        )
        
        # 2D热力图预测头
        self.heatmap_head = nn.Sequential(
            nn.Conv2d(512, 256, 1),
            nn.ReLU(),
            nn.Conv2d(256, num_joints, 1)  # 每个关节一个通道
        )
        
        # 3D深度偏移预测头
        self.depth_head = nn.Sequential(
            nn.Conv2d(512, 256, 1),
            nn.ReLU(),
            nn.Conv2d(256, num_joints, 1)  # 每个关节的深度偏移
        )
        
    def forward(self, depth_img: torch.Tensor) -> dict:
        """
        Args:
            depth_img: (B, 1, H, W) 深度图像, 单位毫米
            
        Returns:
            joints_2d: (B, num_joints, 2) 像素坐标
            joints_3d: (B, num_joints, 3) 3D坐标（相机坐标系）
            confidences: (B, num_joints) 置信度
        """
        B = depth_img.shape[0]
        
        # 特征提取
        feats = self.backbone(depth_img)  # (B, 512, H/8, W/8)
        
        # 2D热力图
        heatmaps = self.heatmap_head(feats)  # (B, num_joints, H/8, W/8)
        heatmaps = F.interpolate(
            heatmaps, scale_factor=8, mode='bilinear', align_corners=False
        )
        
        # 深度偏移
        depth_offsets = self.depth_head(feats)
        depth_offsets = F.interpolate(
            depth_offsets, scale_factor=8, mode='bilinear', align_corners=False
        )
        
        # 提取2D关节点
        joints_2d = self._extract_joints_2d(heatmaps)
        
        # 提取3D关节点
        joints_3d = self._extract_joints_3d(
            joints_2d, depth_img, depth_offsets
        )
        
        # 置信度
        confidences = self._get_confidences(heatmaps)
        
        return {
            'joints_2d': joints_2d,
            'joints_3d': joints_3d,
            'confidences': confidences
        }
    
    def _extract_joints_2d(self, heatmaps: torch.Tensor) -> torch.Tensor:
        """从热力图提取2D坐标"""
        B, J, H, W = heatmaps.shape
        # 每个通道的最大值位置
        heatmaps_flat = heatmaps.view(B, J, -1)
        max_idx = heatmaps_flat.argmax(dim=2)  # (B, J)
        x = (max_idx % W).float()
        y = (max_idx // W).float()
        return torch.stack([x, y], dim=2)  # (B, J, 2)
    
    def _extract_joints_3d(self, joints_2d: torch.Tensor,
                           depth_img: torch.Tensor,
                           depth_offsets: torch.Tensor) -> torch.Tensor:
        """从2D坐标+深度图提取3D坐标"""
        B, J, _ = joints_2d.shape
        _, _, H, W = depth_img.shape
        
        joints_3d = torch.zeros(B, J, 3, device=joints_2d.device)
        
        for b in range(B):
            for j in range(J):
                x, y = joints_2d[b, j].long()
                x = torch.clamp(x, 0, W-1)
                y = torch.clamp(y, 0, H-1)
                
                # 深度值 + 偏移
                z = depth_img[b, 0, y, x] + depth_offsets[b, j, y, x]
                
                # 转换到相机坐标系
                # fx, fy, cx, cy 为相机内参（ToF典型值）
                fx, fy = 320, 320  # 焦距
                cx, cy = 160, 120  # 光心
                
                X = (x - cx) * z / fx
                Y = (y - cy) * z / fy
                Z = z
                
                joints_3d[b, j] = torch.tensor([X, Y, Z], device=joints_2d.device)
        
        return joints_3d
    
    def _get_confidences(self, heatmaps: torch.Tensor) -> torch.Tensor:
        """提取置信度"""
        B, J, H, W = heatmaps.shape
        return heatmaps.view(B, J, -1).max(dim=2)[0]


class STGCNBehaviorRecognition(nn.Module):
    """
    ST-GCN++ 驾驶行为识别
    
    输入：3D关节点序列
    输出：10类行为分类
    
    10类行为定义：
    0: 正常驾驶
    1: 低头看手机
    2: 接打电话
    3: 抽烟
    4: 喝水
    5: 左顾右盼
    6: 调节中控
    7. 身体前倾
    8: 打哈欠/疲劳
    9: 双手离开方向盘
    """
    
    BEHAVIOR_CLASSES = [
        'normal_driving',
        'looking_down_phone',
        'calling',
        'smoking',
        'drinking',
        'looking_around',
        'adjusting_console',
        'leaning_forward',
        'yawning_fatigue',
        'hands_off_wheel'
    ]
    
    def __init__(self, num_joints: int = 16, 
                 num_classes: int = 10,
                 hidden_dim: int = 64,
                 num_layers: int = 3):
        super().__init__()
        
        # 关节空间图卷积
        self.spatial_gcn = SpatialGraphConv(num_joints, hidden_dim)
        
        # 时序卷积
        self.temporal_conv = nn.Sequential(
            nn.Conv1d(hidden_dim, hidden_dim * 2, kernel_size=3, padding=1),
            nn.BatchNorm1d(hidden_dim * 2),
            nn.ReLU(),
            nn.Conv1d(hidden_dim * 2, hidden_dim * 2, kernel_size=3, padding=1),
            nn.BatchNorm1d(hidden_dim * 2),
            nn.ReLU(),
        )
        
        # 多层ST块
        self.st_blocks = nn.ModuleList([
            STBlock(hidden_dim * 2, hidden_dim * 2)
            for _ in range(num_layers)
        ])
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool1d(1),
            nn.Flatten(),
            nn.Linear(hidden_dim * 2, 128),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, num_classes)
        )
        
    def forward(self, joint_sequence: torch.Tensor) -> torch.Tensor:
        """
        Args:
            joint_sequence: (B, T, num_joints, 3) 3D关节点时序
                            T=30 (1秒@30fps)
        Returns:
            logits: (B, num_classes)
        """
        B, T, J, C = joint_sequence.shape
        
        # 空间图卷积
        x = self.spatial_gcn(joint_sequence)  # (B, T, J, hidden_dim*2)
        
        # 重排维度: (B, J*hidden, T)
        x = x.view(B, T, -1).permute(0, 2, 1)  # (B, J*hidden_dim*2, T)
        
        # 时序卷积
        x = self.temporal_conv(x)
        
        # ST块
        for block in self.st_blocks:
            x = block(x)
        
        # 分类
        logits = self.classifier(x)
        return logits


class SpatialGraphConv(nn.Module):
    """空间图卷积：关节点间空间关系建模"""
    
    def __init__(self, num_joints: int, out_dim: int):
        super().__init__()
        # 预定义关节连接关系（16关节）
        # [关节i, 关节j] 表示连接
        self.edges = [
            (0, 1), (1, 2),           # head
            (2, 3), (2, 4),           # neck-shoulder
            (3, 5), (5, 7), (7, 9),   # left arm
            (4, 6), (6, 8), (8, 10),  # right arm  
            (2, 11), (11, 12),        # spine
            (12, 13), (12, 14),       # pelvis-hip
            (13, 15), (14, 16),       # hip-knee
        ]
        
        self.num_joints = num_joints
        self.out_dim = out_dim
        self.linear = nn.Linear(3, out_dim)
        
        # 邻接矩阵
        adj = torch.zeros(num_joints, num_joints)
        for i, j in self.edges:
            adj[i, j] = 1
            adj[j, i] = 1
        adj += torch.eye(num_joints)  # 自环
        # 归一化
        D = adj.sum(dim=1).diag()
        D_inv = D.inverse()
        self.register_buffer('adj', D_inv @ adj)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, T, J, 3)
        Returns:
            (B, T, J, out_dim)
        """
        x = self.linear(x)  # (B, T, J, out_dim)
        # 图卷积
        x = torch.einsum('ij,btjd->btid', self.adj, x)
        return F.relu(x)


class STBlock(nn.Module):
    """ST-GCN 时空块"""
    def __init__(self, in_dim: int, out_dim: int):
        super().__init__()
        self.spatial = nn.Conv1d(in_dim, out_dim, 1)
        self.temporal = nn.Conv1d(out_dim, out_dim, 3, padding=1)
        self.bn = nn.BatchNorm1d(out_dim)
        self.relu = nn.ReLU()
        
    def forward(self, x):
        identity = x
        x = self.spatial(x)
        x = self.temporal(x)
        x = self.bn(x)
        x = self.relu(x + identity)
        return x


# ==================== 测试代码 ====================
if __name__ == "__main__":
    # 测试3D姿态估计
    pose_model = Lightweight3DPoseEstimator(num_joints=16)
    depth_img = torch.randn(2, 1, 240, 320)  # ToF深度图
    with torch.no_grad():
        pose_out = pose_model(depth_img)
    print(f"2D关节点: {pose_out['joints_2d'].shape}")
    print(f"3D关节点: {pose_out['joints_3d'].shape}")
    print(f"置信度: {pose_out['confidences'].shape}")
    print(f"姿态模型参数量: {sum(p.numel() for p in pose_model.parameters())/1e6:.1f}M")
    
    # 测试行为识别
    behav_model = STGCNBehaviorRecognition(num_joints=16, num_classes=10)
    joint_seq = torch.randn(2, 30, 16, 3)  # 30帧
    with torch.no_grad():
        logits = behav_model(joint_seq)
    print(f"\n行为识别输出: {logits.shape}")
    print(f"行为模型参数量: {sum(p.numel() for p in behav_model.parameters())/1e6:.1f}M")
    print(f"姿态+行为总参数量: {(sum(p.numel() for p in pose_model.parameters()) + sum(p.numel() for p in behav_model.parameters()))/1e6:.1f}M")