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

class SimulationDataGenerator:
    """
    阶段1：仿真数据生成
    
    使用 SMPL-X 统计人体模型生成合成深度+IR 图像
    
    流程：
    1. 采样姿态参数 Z (15关节 × 3D)
    2. 采样体型参数 s (身高/体重/比例)
    3. 通过 SMPL-X 生成人体网格
    4. 渲染深度图和 IR 图
    """
    
    def __init__(self, num_joints: int = 15):
        self.num_joints = num_joints
        # SMPL-X 关节层级
        self.kinematic_tree = {
            'pelvis': 0,
            'abdomen': 1, 'thorax': 2, 'neck': 3, 'head': 4,
            'left_hip': 5, 'left_knee': 6,
            'right_hip': 7, 'right_knee': 8,
            'left_shoulder': 9, 'left_elbow': 10, 'left_wrist': 11,
            'right_shoulder': 12, 'right_elbow': 13, 'right_wrist': 14
        }
        
        # 车内典型姿态范围
        self.posture_ranges = {
            'normal_driving': {'neck_pitch': (-10, 10), 'torso_pitch': (-5, 15)},
            'leaning_forward': {'neck_pitch': (20, 45), 'torso_pitch': (20, 40)},
            'leaning_left': {'neck_yaw': (-30, -15), 'torso_roll': (-20, -10)},
            'leaning_right': {'neck_yaw': (15, 30), 'torso_roll': (10, 20)},
            'feet_up': {'hip_flexion': (60, 90), 'knee_flexion': (-30, 0)},
            'sleeping': {'neck_pitch': (30, 60), 'torso_roll': (-15, 15)}
        }
    
    def generate_sample(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        生成一个仿真样本
        
        Returns:
            depth_image: (H, W) 深度图
            ir_image: (H, W) 红外图
            posture_3d: (15, 3) 3D 关节坐标
        """
        # 随机选择姿态类型
        posture_type = np.random.choice(list(self.posture_ranges.keys()))
        ranges = self.posture_ranges[posture_type]
        
        # 采样关节角度
        angles = {}
        for joint, (low, high) in ranges.items():
            angles[joint] = np.random.uniform(low, high)
        
        # 采样体型参数
        height = np.random.uniform(1.5, 1.9)  # 米
        weight = np.random.uniform(50, 100)  # kg
        
        # 生成 3D 关节坐标（简化版前向运动学）
        posture_3d = self._forward_kinematics(angles, height)
        
        # 渲染深度图和 IR 图（模拟）
        H, W = 480, 640
        depth_image = self._render_depth(posture_3d, H, W)
        ir_image = self._render_ir(posture_3d, H, W)
        
        return depth_image, ir_image, posture_3d
    
    def _forward_kinematics(self, angles: dict, height: float) -> np.ndarray:
        """简化版前向运动学"""
        posture = np.zeros((15, 3))
        
        # 骨盆位置（原点）
        posture[0] = [0, 0, 0.4 * height]
        
        # 躯干向上
        posture[1] = posture[0] + [0, 0, 0.15 * height]
        posture[2] = posture[1] + [0, 0, 0.15 * height]
        
        # 颈部
        neck_pitch = np.radians(angles.get('neck_pitch', 0))
        posture[3] = posture[2] + [0, 0.05 * height * np.sin(neck_pitch), 
                                    0.08 * height * np.cos(neck_pitch)]
        
        # 头部
        posture[4] = posture[3] + [0, 0.05 * height * np.sin(neck_pitch), 
                                    0.12 * height * np.cos(neck_pitch)]
        
        # 四肢（简化）
        posture[5] = posture[0] + [0.08, 0, -0.05]  # 左髋
        posture[6] = posture[5] + [0, 0, -0.2]      # 左膝
        posture[7] = posture[0] + [-0.08, 0, -0.05]  # 右髋
        posture[8] = posture[7] + [0, 0, -0.2]       # 右膝
        
        posture[9] = posture[2] + [0.15, 0, 0]   # 左肩
        posture[10] = posture[9] + [0.1, 0, -0.1]  # 左肘
        posture[11] = posture[10] + [0.1, 0, -0.1]  # 左腕
        posture[12] = posture[2] + [-0.15, 0, 0]  # 右肩
        posture[13] = posture[12] + [-0.1, 0, -0.1]
        posture[14] = posture[13] + [-0.1, 0, -0.1]
        
        return posture
    
    def _render_depth(self, posture: np.ndarray, H: int, W: int) -> np.ndarray:
        """模拟深度图渲染"""
        depth = np.full((H, W), 5.0)  # 背景距离 5m
        
        for joint in posture:
            x = int(W/2 + joint[0] * 100)
            y = int(H/2 - joint[2] * 100)
            z = joint[1] + 1.0  # 前方 1m
            
            if 0 <= x < W and 0 <= y < H:
                cv2_circle(depth, (x, y), 15, z, -1)
        
        return depth
    
    def _render_ir(self, posture: np.ndarray, H: int, W: int) -> np.ndarray:
        """模拟 IR 图渲染"""
        ir = np.zeros((H, W), dtype=np.float32)
        
        for joint in posture:
            x = int(W/2 + joint[0] * 100)
            y = int(H/2 - joint[2] * 100)
            
            if 0 <= x < W and 0 <= y < H:
                ir[y-10:y+10, x-10:x+10] = 0.8
        
        return ir


def cv2_circle(img, center, radius, color, thickness):
    """简化版圆形绘制"""
    H, W = img.shape[:2]
    cx, cy = center
    for y in range(max(0, cy-radius), min(H, cy+radius)):
        for x in range(max(0, cx-radius), min(W, cx+radius)):
            if (x-cx)**2 + (y-cy)**2 <= radius**2:
                img[y, x] = color

class OccupantPostureNet(nn.Module):
    """
    车内乘员 3D 姿态估计网络
    
    输入：深度图 + IR 图
    输出：15 个关节的 3D 坐标
    
    架构：双流编码器 + 融合解码器
    """
    
    def __init__(
        self,
        num_joints: int = 15,
        input_size: int = 480,
        hidden_dim: int = 256
    ):
        super().__init__()
        self.num_joints = num_joints
        
        # 深度图编码器
        self.depth_encoder = nn.Sequential(
            nn.Conv2d(1, 32, 3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1))
        )
        
        # IR 图编码器
        self.ir_encoder = nn.Sequential(
            nn.Conv2d(1, 32, 3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1))
        )
        
        # 融合层
        self.fusion = nn.Sequential(
            nn.Linear(256, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.3)
        )
        
        # 关节坐标回归器
        self.regressor = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, num_joints * 3)
        )
        
    def forward(
        self,
        depth: torch.Tensor,  # (B, 1, H, W)
        ir: torch.Tensor      # (B, 1, H, W)
    ) -> torch.Tensor:
        """
        Returns:
            posture_3d: (B, 15, 3) 3D 关节坐标
        """
        # 双流编码
        depth_feat = self.depth_encoder(depth).flatten(1)  # (B, 128)
        ir_feat = self.ir_encoder(ir).flatten(1)            # (B, 128)
        
        # 融合
        fused = torch.cat([depth_feat, ir_feat], dim=-1)    # (B, 256)
        fused = self.fusion(fused)                          # (B, hidden_dim)
        
        # 回归
        posture = self.regressor(fused)                     # (B, 15*3)
        posture = posture.reshape(-1, self.num_joints, 3)   # (B, 15, 3)
        
        return posture


if __name__ == "__main__":
    # 模型初始化
    model = OccupantPostureNet(num_joints=15)
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"模型参数量: {total_params:,}")
    
    # 模拟输入
    depth = torch.randn(4, 1, 480, 640)
    ir = torch.randn(4, 1, 480, 640)
    
    with torch.no_grad():
        posture = model(depth, ir)
    
    print(f"输入: depth={depth.shape}, ir={ir.shape}")
    print(f"输出: posture={posture.shape}")
    
    # 仿真数据生成
    generator = SimulationDataGenerator()
    for i in range(3):
        d, ir_img, pose = generator.generate_sample()
        print(f"\n样本 {i+1}:")
        print(f"  深度图: {d.shape}")
        print(f"  IR图: {ir_img.shape}")
        print(f"  姿态: {pose.shape}")
        print(f"  头部位置: ({pose[4,0]:.2f}, {pose[4,1]:.2f}, {pose[4,2]:.2f})")
    
    print("\n=== 论文报告结果 ===")
    print("中位误差: <10cm (所有关节)")
    print("标注样本: <100 个")
    print("关节点数: 15")
    print("输入: 深度+IR (双 Kinect V2)")