import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class DepthIRFusionNet(nn.Module):
    """
    深度+红外双模态3D姿态估计网络
    
    论文: Tambwekar et al., Sensors 2024
    核心: 轻量级双分支特征提取 + 深度引导的3D坐标恢复
    """
    
    def __init__(
        self,
        num_joints: int = 14,
        depth_channels: int = 1,
        ir_channels: int = 1,
        hidden_dim: int = 128
    ):
        super().__init__()
        self.num_joints = num_joints
        
        # 深度分支 - 提取深度结构特征
        self.depth_encoder = nn.Sequential(
            nn.Conv2d(depth_channels, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((16, 16))
        )
        
        # 红外分支 - 提取纹理/轮廓特征
        self.ir_encoder = nn.Sequential(
            nn.Conv2d(ir_channels, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((16, 16))
        )
        
        # 融合层
        self.fusion = nn.Sequential(
            nn.Conv2d(256, 256, 1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.Conv2d(256, 128, 1),
            nn.BatchNorm2d(128),
            nn.ReLU()
        )
        
        # 2D关键点检测头
        self.heatmap_head = nn.Sequential(
            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, num_joints, 1)
        )
        
        # 3D坐标回归头
        self.regression_head = nn.Sequential(
            nn.Linear(128 * 16 * 16, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, num_joints * 3)
        )
    
    def forward(
        self, 
        depth: torch.Tensor, 
        ir: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            depth: 深度图, shape=(B, 1, H, W), 单位mm
            ir: 红外图, shape=(B, 1, H, W), 0-255
            
        Returns:
            heatmaps: 2D关键点热图, shape=(B, num_joints, H/4, W/4)
            joints_3d: 3D关节坐标, shape=(B, num_joints, 3), 单位mm
        """
        # 双分支特征提取
        depth_feat = self.depth_encoder(depth)    # (B, 128, 16, 16)
        ir_feat = self.ir_encoder(ir)               # (B, 128, 16, 16)
        
        # 特征融合
        fused = torch.cat([depth_feat, ir_feat], dim=1)  # (B, 256, 16, 16)
        fused = self.fusion(fused)                         # (B, 128, 16, 16)
        
        # 2D热图预测
        heatmaps = self.heatmap_head(fused)  # (B, num_joints, H/4, W/4)
        
        # 3D坐标回归
        flat = fused.flatten(1)               # (B, 128*16*16)
        joints_3d = self.regression_head(flat) # (B, num_joints*3)
        joints_3d = joints_3d.reshape(-1, self.num_joints, 3)
        
        return heatmaps, joints_3d


def depth_to_3d_joints(
    heatmap: np.ndarray,
    depth_map: np.ndarray,
    camera_intrinsics: dict
) -> np.ndarray:
    """
    从2D热图+深度图恢复3D关节坐标
    
    Args:
        heatmap: 关键点热图, shape=(num_joints, H, W)
        depth_map: 深度图, shape=(H, W), 单位mm
        camera_intrinsics: 相机内参
            fx, fy: 焦距
            cx, cy: 主点
    
    Returns:
        joints_3d: 3D关节坐标, shape=(num_joints, 3), 单位mm
    """
    fx = camera_intrinsics['fx']
    fy = camera_intrinsics['fy']
    cx = camera_intrinsics['cx']
    cy = camera_intrinsics['cy']
    
    num_joints = heatmap.shape[0]
    joints_3d = np.zeros((num_joints, 3))
    
    for j in range(num_joints):
        # 从热图找最大值位置
        v, u = np.unravel_index(np.argmax(heatmap[j]), heatmap[j].shape)
        
        # 下采样补偿（热图分辨率 < 原图）
        scale_h = depth_map.shape[0] / heatmap[j].shape[0]
        scale_w = depth_map.shape[1] / heatmap[j].shape[1]
        u_full = int(u * scale_w)
        v_full = int(v * scale_h)
        
        # 获取深度值
        z = depth_map[v_full, u_full]
        
        # 像素坐标 → 相机坐标
        x = (u_full - cx) * z / fx
        y = (v_full - cy) * z / fy
        
        joints_3d[j] = [x, y, z]
    
    return joints_3d


# 完整测试
if __name__ == "__main__":
    # 相机内参 (Intel RealSense D435i)
    intrinsics = {
        'fx': 615.0, 'fy': 615.0,
        'cx': 320.0, 'cy': 240.0
    }
    
    model = DepthIRFusionNet(num_joints=14)
    
    # 模拟输入
    depth = torch.randn(2, 1, 480, 640) * 1000  # mm
    ir = torch.rand(2, 1, 480, 640) * 255        # 0-255
    
    heatmaps, joints_3d = model(depth, ir)
    print(f"热图: {heatmaps.shape}")
    print(f"3D关节: {joints_3d.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")
    
    # 3D坐标恢复测试
    heatmaps_np = heatmaps[0].detach().numpy()
    depth_np = depth[0, 0].numpy()
    joints = depth_to_3d_joints(heatmaps_np, depth_np, intrinsics)
    print(f"恢复的3D坐标 (mm): {joints[0]}")

class FewShotTrainer:
    """
    少样本3D姿态估计训练器
    
    策略:
    1. 预训练于公开3D姿态数据集（如Human3.6M）
    2. 域适应: 用少量车内标注微调
    3. 数据增强: 深度图扰动 + 红外图噪声模拟
    """
    
    def __init__(self, model: nn.Module, lr: float = 1e-4):
        self.model = model
        self.optimizer = torch.optim.AdamW(
            model.parameters(), lr=lr, weight_decay=1e-4
        )
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer, T_max=100
        )
    
    def augment_depth(self, depth: torch.Tensor) -> torch.Tensor:
        """深度图增强: 添加模拟噪声"""
        noise = torch.randn_like(depth) * 10  # 10mm噪声
        return depth + noise
    
    def augment_ir(self, ir: torch.Tensor) -> torch.Tensor:
        """红外图增强: 亮度抖动 + 模拟遮挡"""
        # 亮度调整
        brightness = 0.8 + torch.rand(1).item() * 0.4
        ir = ir * brightness
        # 随机遮挡
        mask = torch.rand_like(ir) > 0.95
        ir[mask] = 0
        return ir.clamp(0, 255)
    
    def train_step(
        self, 
        depth: torch.Tensor, 
        ir: torch.Tensor,
        gt_heatmaps: torch.Tensor,
        gt_joints_3d: torch.Tensor
    ) -> dict:
        """单步训练"""
        # 数据增强
        depth_aug = self.augment_depth(depth)
        ir_aug = self.augment_ir(ir)
        
        # 前向
        pred_heatmaps, pred_joints = self.model(depth_aug, ir_aug)
        
        # 损失
        loss_heatmap = F.mse_loss(pred_heatmaps, gt_heatmaps)
        loss_joints = F.l1_loss(pred_joints, gt_joints_3d)
        loss = loss_heatmap + 0.5 * loss_joints
        
        # 反向
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        
        return {
            'loss': loss.item(),
            'heatmap_loss': loss_heatmap.item(),
            'joint_loss': loss_joints.item()
        }

def classify_oop_posture(joints_3d: np.ndarray) -> dict:
    """
    基于OOP（异常姿态）分类
    
    Args:
        joints_3d: 14个关节3D坐标, shape=(14, 3), 单位mm
    
    Returns:
        result: 分类结果 + 风险等级
    """
    # 关节索引定义 (14 joints)
    # 0:头 1:颈 2:左肩 3:右肩 4:左肘 5:右肘
    # 6:左髋 7:右髋 8:左膝 9:右膝 10:躯干中心
    
    head = joints_3d[0]
    neck = joints_3d[1]
    torso_center = joints_3d[10]
    left_hip = joints_3d[6]
    right_hip = joints_3d[7]
    hip_center = (left_hip + right_hip) / 2
    
    # 计算关键角度
    # 头-躯干角度（检测前倾/后仰）
    head_torso_vec = head - hip_center
    head_torso_angle = np.degrees(
        np.arctan2(head_torso_vec[0], head_torso_vec[2])
    )
    
    # 躯干倾斜角（检测侧倾）
    torso_vec = neck - hip_center
    torso_lateral_angle = np.degrees(
        np.arctan2(torso_vec[0], torso_vec[1])
    )
    
    # 分类逻辑
    posture = "normal"
    risk_level = "safe"
    
    if abs(head_torso_angle) > 30:
        posture = "forward_lean" if head_torso_angle > 0 else "backward_lean"
        risk_level = "high_risk"
    elif abs(torso_lateral_angle) > 25:
        posture = "lateral_lean"
        risk_level = "medium_risk"
    
    return {
        'posture': posture,
        'risk_level': risk_level,
        'head_torso_angle': head_torso_angle,
        'torso_lateral_angle': torso_lateral_angle,
        'meets_ncap': abs(head_torso_angle) < 15  # NCAP正常坐姿阈值
    }


# 测试
if __name__ == "__main__":
    # 模拟正常坐姿
    normal_joints = np.array([
        [0, -200, 500],   # 头
        [0, -150, 500],   # 颈
        [-150, -150, 500], # 左肩
        [150, -150, 500],  # 右肩
        # ... 简化
        [0, 0, 500],      # 躯干中心
        [-100, 0, 500],   # 左髋
        [100, 0, 500],    # 右髋
    ])
    
    result = classify_oop_posture(normal_joints)
    print(f"姿态: {result['posture']}, 风险: {result['risk_level']}")