import torch
import torch.nn as nn

class PressureNet(nn.Module):
    """
    PressureNet: 从压力图像预测 3D 人体姿态和体型
    
    论文 Clever et al. (2020) 的简化实现
    
    输入: 压力分布图 (B, 1, H, W)
    输出: SMPL-X 参数 (pose: 72维, shape: 10维)
    """
    
    def __init__(self, pressure_size=(72, 36)):
        super().__init__()
        
        # 编码器：从压力图提取特征
        self.encoder = nn.Sequential(
            # Conv Block 1
            nn.Conv2d(1, 32, 5, stride=2, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            
            # Conv Block 2
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            
            # Conv Block 3
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            
            # Conv Block 4
            nn.Conv2d(128, 256, 3, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
        )
        
        # 计算 flatten 后的尺寸
        h, w = pressure_size
        flat_h, flat_w = h // 16, w // 16
        if flat_h == 0: flat_h = 1
        if flat_w == 0: flat_w = 1
        flat_dim = 256 * max(flat_h, 1) * max(flat_w, 1)
        
        # 全连接层
        self.fc = nn.Sequential(
            nn.Linear(flat_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
        )
        
        # 姿态回归头（SMPL pose 参数: 24关节 × 3 = 72）
        self.pose_head = nn.Linear(256, 72)
        
        # 体型回归头（SMPL shape 参数: 10）
        self.shape_head = nn.Linear(256, 10)
        
        # 全局平移（3D）
        self.trans_head = nn.Linear(256, 3)
        
    def forward(self, pressure_img):
        """
        Args:
            pressure_img: (B, 1, H, W) 归一化压力图
            
        Returns:
            dict: pose, shape, translation
        """
        feat = self.encoder(pressure_img)
        feat = feat.flatten(1)
        feat = self.fc(feat)
        
        pose = self.pose_head(feat)
        shape = self.shape_head(feat)
        trans = self.trans_head(feat)
        
        return {
            'pose': pose,      # (B, 72)
            'shape': shape,     # (B, 10)
            'trans': trans,     # (B, 3)
        }


class PressureSynthesizer:
    """
    合成压力图像生成器
    
    论文核心：用物理仿真生成训练数据
    """
    
    def __init__(self, grid_size=(72, 36)):
        self.grid_h, self.grid_w = grid_size
        self.cell_size = 0.027  # 2.7cm 间距
    
    def generate(self, pose_params, shape_params, body_mass=70):
        """
        从 SMPL 参数生成合成压力图
        
        Args:
            pose_params: 24x3 关节角度
            shape_params: 10 体型参数
            body_mass: 体重 (kg)
            
        Returns:
            pressure_map: (H, W) 压力分布
        """
        # 简化版：基于高斯分布模拟接触压力
        pressure_map = np.zeros((self.grid_h, self.grid_w))
        
        # 主要接触区域：头部、背部、臀部、脚
        # 根据姿态参数确定接触位置
        contact_regions = self._get_contact_regions(pose_params)
        
        for region in contact_regions:
            cx, cy, intensity = region
            # 高斯分布模拟压力
            for i in range(self.grid_h):
                for j in range(self.grid_w):
                    dist = np.sqrt((i - cx)**2 + (j - cy)**2)
                    pressure_map[i, j] += intensity * np.exp(-dist**2 / 4)
        
        # 归一化
        pressure_map = pressure_map / (pressure_map.max() + 1e-6)
        
        return pressure_map
    
    def _get_contact_regions(self, pose_params):
        """根据姿态参数确定接触区域"""
        # 简化：返回接触点的坐标和压力强度
        return [
            (10, 18, 0.3),   # 头部
            (30, 18, 0.5),   # 上背
            (45, 18, 0.8),   # 臀部
            (60, 18, 0.4),   # 小腿
            (65, 18, 0.3),   # 脚
        ]


import numpy as np

# 测试
if __name__ == "__main__":
    # 生成合成压力图
    synth = PressureSynthesizer()
    pose = np.zeros((24, 3))
    shape = np.zeros(10)
    pressure_img = synth.generate(pose, shape)
    
    print(f"压力图尺寸: {pressure_img.shape}")
    print(f"压力值范围: {pressure_img.min():.3f} - {pressure_img.max():.3f}")
    print(f"非零点数: {np.count_nonzero(pressure_img > 0.01)}")
    
    # PressureNet 推理
    model = PressureNet()
    x = torch.FloatTensor(pressure_img).unsqueeze(0).unsqueeze(0)
    output = model(x)
    
    print(f"\nPressureNet 输出:")
    print(f"  Pose: {output['pose'].shape}")
    print(f"  Shape: {output['shape'].shape}")
    print(f"  Trans: {output['trans'].shape}")
    print(f"  参数总数: {sum(p.numel() for p in model.parameters()):,}")

class CabinOOPDetector:
    """
    座舱 OOP 检测器：压力垫 + 摄像头融合
    
    压力垫提供：接触分布、重心位置、左右不对称
    摄像头提供：上半身姿态、头部位置、安全带状态
    """
    
    def __init__(self):
        self.pressure_model = PressureNet(pressure_size=(32, 32))
        # 加载预训练权重（从床压迁移）
        self.confidence_threshold = 0.7
    
    def detect(self, pressure_img, camera_feat=None):
        """
        OOP 检测
        
        Args:
            pressure_img: (32, 32) 座椅压力图
            camera_feat: 可选的摄像头特征
            
        Returns:
            oop_result: 姿态分类和置信度
        """
        # 1. 压力垫推理
        x = torch.FloatTensor(pressure_img).unsqueeze(0).unsqueeze(0)
        pressure_output = self.pressure_model(x)
        
        # 2. 提取压力特征
        pressure_feat = self._extract_pressure_features(pressure_img)
        
        # 3. 分类
        posture = self._classify_posture(pressure_feat)
        
        # 4. 如果有摄像头特征，融合提升
        if camera_feat is not None:
            posture = self._fusion_refine(posture, pressure_feat, camera_feat)
        
        return {
            'posture': posture,
            'confidence': pressure_feat['confidence'],
            'pressure_features': pressure_feat,
            'is_oop': posture not in ['normal_sitting'],
        }
    
    def _extract_pressure_features(self, img):
        """从压力图提取特征"""
        total = img.sum()
        if total < 0.1:
            return {'confidence': 0.0, 'empty': True}
        
        # 重心
        ys, xs = np.mgrid[0:img.shape[0], 0:img.shape[1]]
        cx = (xs * img).sum() / total
        cy = (ys * img).sum() / total
        
        # 左右不对称
        left = img[:, :img.shape[1]//2].sum()
        right = img[:, img.shape[1]//2:].sum()
        asymmetry = abs(left - right) / (left + right + 1e-6)
        
        # 前后分布（背=上半，臀=下半）
        upper = img[:img.shape[0]//2, :].sum()
        lower = img[img.shape[0]//2:, :].sum()
        front_back_ratio = lower / (upper + 1e-6)
        
        return {
            'confidence': min(total / 10, 1.0),
            'center': (cx, cy),
            'asymmetry': asymmetry,
            'front_back_ratio': front_back_ratio,
            'total_pressure': total,
            'empty': False
        }
    
    def _classify_posture(self, feat):
        """基于规则+ML的姿态分类"""
        if feat.get('empty', True):
            return 'empty'
        
        if feat['asymmetry'] > 0.3:
            return 'leaning_side'
        
        if feat['front_back_ratio'] > 2.5:
            return 'slouching'  # 瘫坐
        
        if feat['front_back_ratio'] < 1.0:
            return 'leaning_forward'  # 前倾
        
        return 'normal_sitting'


# 测试
if __name__ == "__main__":
    detector = CabinOOPDetector()
    
    # 场景1: 正常坐姿
    pressure_normal = np.zeros((32, 32))
    pressure_normal[8:16, 8:24] = 0.6  # 臀部
    pressure_normal[2:8, 10:22] = 0.4  # 背部
    r1 = detector.detect(pressure_normal)
    print(f"正常: {r1['posture']}, OOP={r1['is_oop']}")
    
    # 场景2: 前倾
    pressure_lean = np.zeros((32, 32))
    pressure_lean[20:28, 8:24] = 0.7  # 大腿（前倾重心前移）
    pressure_lean[4:8, 10:22] = 0.2   # 背部（压力减少）
    r2 = detector.detect(pressure_lean)
    print(f"前倾: {r2['posture']}, OOP={r2['is_oop']}")
    
    # 场景3: 侧倾
    pressure_side = np.zeros((32, 32))
    pressure_side[8:16, 4:14] = 0.7  # 左侧
    pressure_side[8:16, 14:24] = 0.2 # 右侧
    pressure_side[2:8, 6:14] = 0.3
    r3 = detector.detect(pressure_side)
    print(f"侧倾: {r3['posture']}, OOP={r3['is_oop']}")