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

class PressureNet(nn.Module):
    """
    PressureNet: 从压力图像估计3D人体姿态和体型
    
    两阶段架构：
    - Mod1: 粗略估计
    - Mod2: 精细优化
    
    关键创新：Pressure Map Reconstruction (PMR)
    - 从估计的3D mesh重建压力图
    - 与输入压力图对比，确保一致性
    
    Args:
        pressure_size: 压力垫分辨率 (H, W)
        n_mesh_vertices: SMPL mesh 顶点数
        gender_dim: 性别编码维度
    """
    
    def __init__(self, pressure_h=64, pressure_w=32, 
                 n_mesh_vertices=6890, gender_dim=2):
        super().__init__()
        
        # === Mod1: 粗略估计 ===
        self.mod1 = PressureEstimator(
            pressure_h, pressure_w, n_mesh_vertices, gender_dim,
            use_reconstruction=True  # PMR
        )
        
        # === Mod2: 精细优化 ===
        self.mod2 = PressureEstimator(
            pressure_h, pressure_w, n_mesh_vertices, gender_dim,
            use_reconstruction=True,
            input_channels=2  # 原始压力图 + Mod1重建图
        )
        
    def forward(self, pressure_image, gender):
        """
        Args:
            pressure_image: (B, 1, H, W) 压力分布图
            gender: (B, 2) one-hot性别编码
        
        Returns:
            mesh_vertices: (B, V, 3) 3D mesh顶点坐标
            reconstructed_pressure: (B, 1, H, W) 重建压力图
        """
        # Mod1: 粗略估计
        mesh1, recon1 = self.mod1(pressure_image, gender)
        
        # Mod2: 精细优化（输入：原始+Mod1重建）
        mod2_input = torch.cat([pressure_image, recon1], dim=1)
        mesh2, recon2 = self.mod2(mod2_input, gender)
        
        return mesh2, recon2


class PressureEstimator(nn.Module):
    """单阶段压力→3D mesh估计器"""
    
    def __init__(self, pressure_h, pressure_w, n_vertices, 
                 gender_dim, use_reconstruction=True, input_channels=1):
        super().__init__()
        self.use_reconstruction = use_reconstruction
        
        # CNN编码器
        self.encoder = nn.Sequential(
            nn.Conv2d(input_channels, 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.MaxPool2d(2),
            
            nn.Conv2d(128, 256, 3, padding=1),
            nn.BatchNorm2d(256), nn.ReLU(),
            nn.MaxPool2d(2),
            
            nn.Conv2d(256, 512, 3, padding=1),
            nn.BatchNorm2d(512), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        
        # 性别编码
        self.gender_encoder = nn.Linear(gender_dim, 64)
        
        # SMPL参数预测
        # pose: 72维 (24关节×3轴), shape: 10维
        self.pose_head = nn.Sequential(
            nn.Linear(512 + 64, 256),
            nn.ReLU(),
            nn.Linear(256, 72)  # SMPL pose
        )
        self.shape_head = nn.Sequential(
            nn.Linear(512 + 64, 256),
            nn.ReLU(),
            nn.Linear(256, 10)  # SMPL shape
        )
        
        # Mesh回归（从SMPL参数到顶点）
        self.mesh_regressor = SMPLMeshRegressor(n_vertices)
        
        # PMR: 压力图重建
        if use_reconstruction:
            self.pmr = PressureMapReconstructor(
                n_vertices, pressure_h, pressure_w
            )
    
    def forward(self, pressure_image, gender):
        # CNN编码
        feat = self.encoder(pressure_image).flatten(1)  # (B, 512)
        gender_feat = self.gender_encoder(gender)  # (B, 64)
        combined = torch.cat([feat, gender_feat], dim=1)
        
        # SMPL参数
        pose = self.pose_head(combined)
        shape = self.shape_head(combined)
        
        # Mesh回归
        mesh = self.mesh_regressor(pose, shape)  # (B, V, 3)
        
        # PMR重建
        recon_pressure = None
        if self.use_reconstruction:
            recon_pressure = self.pmr(mesh, pressure_image.shape[-2:])
        
        return mesh, recon_pressure


class PressureMapReconstructor(nn.Module):
    """PMR: 从3D mesh重建压力图"""
    
    def __init__(self, n_vertices, out_h, out_w):
        super().__init__()
        self.out_h = out_h
        self.out_w = out_w
        
        # 将3D顶点投影到压力垫平面
        self.proj = nn.Sequential(
            nn.Linear(n_vertices * 3, 1024),
            nn.ReLU(),
            nn.Linear(1024, out_h * out_w)
        )
        
    def forward(self, mesh, target_size):
        B, V, _ = mesh.shape
        flat = mesh.view(B, -1)  # (B, V*3)
        pressure = self.proj(flat)  # (B, H*W)
        pressure = pressure.view(B, 1, self.out_h, self.out_w)
        return F.relu(pressure)  # 压力非负


class SMPLMeshRegressor(nn.Module):
    """简化的SMPL mesh回归器"""
    def __init__(self, n_vertices):
        super().__init__()
        self.regressor = nn.Sequential(
            nn.Linear(72 + 10, 512),  # pose + shape
            nn.ReLU(),
            nn.Linear(512, 1024),
            nn.ReLU(),
            nn.Linear(1024, n_vertices * 3)
        )
        
    def forward(self, pose, shape):
        x = torch.cat([pose, shape], dim=1)
        mesh = self.regressor(x)
        return mesh.view(-1, 6890, 3)  # (B, V, 3)


# === IMS 座椅OOP检测集成 ===
class SeatPressureOOPDetector:
    """座椅压力传感器OOP检测"""
    
    OOP_CATEGORIES = {
        0: 'normal',           # 正常坐姿
        1: 'forward_lean',     # 前倾
        2: 'reclined',         # 后仰
        3: 'side_lean',        # 侧倾
        4: 'child_seat',       # 儿童座椅
        5: 'empty',            # 空座
        6: 'feet_on_dash',     # 脚搭仪表台（需摄像头辅助）
    }
    
    def __init__(self, sensor_h=32, sensor_w=16):
        self.sensor_h = sensor_h  # 座椅压力阵列行数
        self.sensor_w = sensor_w  # 列数
        self.model = PressureNet(
            pressure_h=sensor_h, pressure_w=sensor_w
        )
        
    def detect_oop(self, pressure_image):
        """OOP检测"""
        # 估计3D姿态
        mesh, recon = self.model(pressure_image, torch.tensor([[1, 0]]))
        
        # 从mesh提取姿态特征
        posture_features = self._extract_posture_features(mesh)
        
        # OOP分类
        oop_class = self._classify_oop(posture_features, pressure_image)
        
        return {
            'oop_category': self.OOP_CATEGORIES[oop_class],
            'confidence': 0.85,
            'mesh': mesh,
            'posture_features': posture_features
        }
    
    def _extract_posture_features(self, mesh):
        """从3D mesh提取姿态特征"""
        # 躯干角度
        # 头部位置
        # 四肢位置
        return {'torso_angle': 15, 'head_height': 0.8}
    
    def _classify_oop(self, features, pressure):
        """OOP分类"""
        total_pressure = pressure.sum()
        if total_pressure < 0.1:  # 空座
            return 5
        # 简化分类
        return 0  # normal


# 测试
if __name__ == "__main__":
    B = 4
    pressure = torch.randn(B, 1, 64, 32)
    gender = torch.tensor([[1, 0]] * B, dtype=torch.float32)
    
    model = PressureNet(64, 32, 6890)
    mesh, recon = model(pressure, gender)
    
    print(f"压力输入: {pressure.shape}")
    print(f"3D Mesh: {mesh.shape}")
    print(f"重建压力: {recon.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters())/1e6:.2f}M")
    
    # OOP检测
    detector = SeatPressureOOPDetector(32, 16)
    result = detector.detect_oop(pressure[:1])
    print(f"\nOOP检测结果: {result['oop_category']}")