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

class PressureMapReconstruction(nn.Module):
    """
    压力图重建网络（PMR）
    
    核心创新：从估计的3D人体Mesh重建压力图，
    与输入压力图对比，强制一致性。
    
    论文Section 3.3: Pressure Map Reconstruction
    
    Args:
        mesh_dim: 人体Mesh顶点特征维度
        grid_size: 压力传感器阵列尺寸
    """
    
    def __init__(self, mesh_dim: int = 6890, grid_size: tuple = (64, 27)):
        super().__init__()
        self.grid_h, self.grid_w = grid_size
        
        # Mesh特征 → 压力图
        self.decoder = nn.Sequential(
            nn.Linear(mesh_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 1024),
            nn.ReLU(),
            nn.Linear(1024, self.grid_h * self.grid_w),
            nn.Sigmoid()  # 压力值0-1
        )
        
    def forward(self, mesh_vertices: torch.Tensor) -> torch.Tensor:
        """
        Args:
            mesh_vertices: (B, mesh_dim) 人体Mesh顶点坐标
            
        Returns:
            pressure_map: (B, 1, grid_h, grid_w) 重建的压力图
        """
        pressure = self.decoder(mesh_vertices)
        return pressure.view(-1, 1, self.grid_h, self.grid_w)


class PressureNetModule(nn.Module):
    """
    PressureNet单模块
    
    架构: [压力图] → [编码器] → [3D Mesh] → [PMR验证] → [输出]
    
    论文Section 3.2: PressureNet Architecture
    """
    
    def __init__(self, grid_size=(64, 27), mesh_dim=6890, 
                 use_pmr: bool = True):
        super().__init__()
        self.use_pmr = use_pmr
        h, w = grid_size
        
        # 压力图编码器（类似hourglass）
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),  # (32, h/2, w/2)
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),  # (64, h/4, w/4)
            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),  # (128, 1, 1)
        )
        
        # Mesh预测头
        self.mesh_head = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128, 512),
            nn.ReLU(),
            nn.Linear(512, 1024),
            nn.ReLU(),
            nn.Linear(1024, mesh_dim * 3)  # xyz坐标
        )
        
        # PMR网络
        if use_pmr:
            self.pmr = PressureMapReconstruction(mesh_dim * 3, grid_size)
    
    def forward(self, pressure_map: torch.Tensor):
        """
        Args:
            pressure_map: (B, 1, H, W) 压力传感垫读数
            
        Returns:
            mesh: (B, mesh_dim, 3) 3D人体Mesh
            reconstructed_pressure: (B, 1, H, W) 重建压力图
        """
        # 编码
        feat = self.encoder(pressure_map)
        
        # 预测Mesh
        mesh = self.mesh_head(feat).view(-1, 6890, 3)
        
        # PMR一致性检查
        if self.use_pmr:
            recon_pressure = self.pmr(mesh.flatten(1))
        else:
            recon_pressure = None
            
        return mesh, recon_pressure


class PressureNet(nn.Module):
    """
    PressureNet: 两阶段3D人体姿态估计
    
    Mod1: 粗估计
    Mod2: 精细化（利用Mod1的重建压力图作为额外输入）
    
    论文Section 3.2: Two-Stage Architecture
    """
    
    def __init__(self, grid_size=(64, 27), mesh_dim=6890):
        super().__init__()
        # Mod1: 粗估计
        self.mod1 = PressureNetModule(grid_size, mesh_dim, use_pmr=True)
        # Mod2: 精细化
        self.mod2 = PressureNetModule(grid_size, mesh_dim, use_pmr=True)
        
    def forward(self, pressure_map: torch.Tensor):
        """
        Args:
            pressure_map: (B, 1, 64, 27)
            
        Returns:
            mesh_final: (B, 6890, 3) 最终Mesh
            losses: dict of intermediate losses
        """
        # Stage 1: 粗估计
        mesh_coarse, recon1 = self.mod1(pressure_map)
        
        # Stage 2: 精细化（输入 = 原始 + Mod1重建）
        if recon1 is not None:
            enhanced_input = torch.cat([pressure_map, recon1], dim=1)
        else:
            enhanced_input = pressure_map
        
        mesh_final, recon2 = self.mod2(enhanced_input)
        
        return mesh_final, {'recon1': recon1, 'recon2': recon2}


# ============ 完整测试代码 ============
if __name__ == "__main__":
    import numpy as np
    import time
    
    # PressurePose数据集参数
    GRID_H, GRID_W = 64, 27  # 压力传感阵列尺寸
    MESH_DIM = 6890           # SMPL Mesh顶点数
    BATCH_SIZE = 8
    
    # 初始化模型
    model = PressureNet(
        grid_size=(GRID_H, GRID_W),
        mesh_dim=MESH_DIM
    )
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"PressureNet参数量: {total_params:,}")
    
    # 模拟压力图数据
    np.random.seed(42)
    
    def generate_pressure_image(pose_type: str, batch=1):
        """
        生成模拟压力传感垫图像
        
        pose_type:
          - 'supine': 仰卧
          - 'side': 侧卧
          - 'reclined': 半躺（OOP场景）
          - 'empty': 空座
        """
        img = np.zeros((batch, 1, GRID_H, GRID_W), dtype=np.float32)
        
        if pose_type == 'supine':
            # 仰卧：均匀分布的中部高压力
            img[:, 0, 20:45, 8:20] = 0.6 + np.random.randn() * 0.1
            img[:, 0, 15:50, 5:22] += 0.3
        elif pose_type == 'side':
            # 侧卧：偏移的压力分布
            img[:, 0, 25:50, 3:12] = 0.8
            img[:, 0, 30:45, 2:10] = 0.9
        elif pose_type == 'reclined':
            # 半躺（OOP）：不均匀分布
            img[:, 0, 10:30, 15:25] = 0.7  # 上半身
            img[:, 0, 35:50, 8:18] = 0.4   # 下半身
            img[:, 0, 5:15, 18:27] = 0.5   # 头部偏移
        elif pose_type == 'empty':
            # 空座：极低噪声
            img += np.random.randn(batch, 1, GRID_H, GRID_W).astype(np.float32) * 0.02
        
        # 加噪声
        img += np.random.randn(batch, 1, GRID_H, GRID_W).astype(np.float32) * 0.05
        return np.clip(img, 0, 1)
    
    # 测试不同姿态
    print("\n--- 压力图生成测试 ---")
    poses = ['supine', 'side', 'reclined', 'empty']
    for pose in poses:
        pressure = generate_pressure_image(pose, batch=1)
        print(f"  {pose:10s}: mean={pressure.mean():.4f}, "
              f"sum={pressure.sum():.4f}, "
              f"max={pressure.max():.4f}")
    
    # 前向传播
    print("\n--- 前向传播 ---")
    model.eval()
    
    # 生成batch数据
    x_batch = torch.cat([
        torch.tensor(generate_pressure_image('supine', 2)),
        torch.tensor(generate_pressure_image('side', 2)),
        torch.tensor(generate_pressure_image('reclined', 2)),
        torch.tensor(generate_pressure_image('empty', 2)),
    ])
    
    with torch.no_grad():
        mesh, aux = model(x_batch)
        print(f"输入压力图: {x_batch.shape}")
        print(f"输出Mesh: {mesh.shape}")
        print(f"重建压力图1: {aux['recon1'].shape}")
        
        # 压力图重建误差
        recon_error = F.mse_loss(aux['recon1'], x_batch)
        print(f"重建MSE: {recon_error.item():.6f}")
    
    # 训练流程
    print("\n--- 训练流程 ---")
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    
    model.train()
    for epoch in range(3):
        optimizer.zero_grad()
        
        # 前向
        mesh, aux = model(x_batch)
        
        # 多任务损失
        # 1. PMR重建损失
        recon_loss = criterion(aux['recon1'], x_batch)
        recon_loss2 = criterion(aux['recon2'], x_batch)
        
        # 2. Mesh正则化（平滑性）
        mesh_reg = torch.mean(torch.abs(mesh[:, 1:] - mesh[:, :-1]))
        
        loss = recon_loss + 0.5 * recon_loss2 + 0.01 * mesh_reg
        
        loss.backward()
        optimizer.step()
        
        print(f"Epoch {epoch+1}: Total={loss.item():.4f}, "
              f"Recon1={recon_loss.item():.4f}, "
              f"Recon2={recon_loss2.item():.4f}")
    
    # 推理速度
    model.eval()
    with torch.no_grad():
        start = time.time()
        for _ in range(50):
            _ = model(x_batch[:1])
        elapsed = (time.time() - start) / 50
        print(f"\n单样本推理延迟: {elapsed*1000:.1f}ms")
    
    # OOP检测应用
    print("\n--- OOP检测应用 ---")
    
    def detect_oop(pressure_map, model, threshold=0.3):
        """
        基于压力图检测OOP异常姿态
        
        Returns:
            dict: {
                'is_oop': bool,
                'pressure_distribution': str,
                'contact_area': float,
                'confidence': float
            }
        """
        with torch.no_grad():
            mesh, aux = model(pressure_map.unsqueeze(0))
            recon = aux['recon1']
            
            # 重建误差大 → 异常姿态
            recon_error = F.mse_loss(recon, pressure_map.unsqueeze(0)).item()
            
            # 压力分布分析
            pressure = pressure_map.squeeze().numpy()
            total_pressure = pressure.sum()
            if total_pressure < 1.0:
                dist = 'empty'
            else:
                # 重心位置
                y_coords = np.arange(GRID_H)
                x_coords = np.arange(GRID_W)
                cy = np.sum(pressure * y_coords[:, None]) / max(total_pressure, 1e-6)
                cx = np.sum(pressure * x_coords[None, :]) / max(total_pressure, 1e-6)
                
                if 25 < cy < 40 and 8 < cx < 18:
                    dist = 'normal_seated'
                elif cy < 20:
                    dist = 'reclined_oop'
                elif cx < 5 or cx > 20:
                    dist = 'side_oop'
                else:
                    dist = 'unknown_oop'
            
            contact_area = (pressure > 0.1).sum() / (GRID_H * GRID_W)
            
        return {
            'is_oop': recon_error > threshold or 'oop' in dist,
            'pressure_distribution': dist,
            'contact_area': float(contact_area),
            'confidence': 1.0 - min(recon_error * 5, 0.5)
        }
    
    # 测试OOP检测
    for pose in poses:
        pressure = torch.tensor(generate_pressure_image(pose, 1)).squeeze(0)
        result = detect_oop(pressure, model)
        print(f"  {pose:10s}: OOP={result['is_oop']}, "
              f"dist={result['pressure_distribution']}, "
              f"area={result['contact_area']:.2%}, "
              f"conf={result['confidence']:.2f}")

class PressureBasedOMS:
    """
    基于压力垫的OMS系统
    
    硬件:
      - 压力传感垫: 嵌入座椅面套
      - 采集芯片: MCU + ADC (ESP32-S3级别)
      - 通信: CAN-FD到域控
    
    软件:
      - PressureNet: INT8量化, <2MB
      - OOP分类器: 基于压力分布特征
      - 融合接口: 与摄像头/UWB联合决策
    """
    
    def __init__(self):
        self.model = PressureNet(grid_size=(64, 27))
        self.oop_threshold = 0.3
        self.history = []
        
    def process_frame(self, pressure_raw: np.ndarray) -> dict:
        """处理一帧压力数据"""
        # 预处理
        pressure = torch.tensor(pressure_raw, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
        
        # OOP检测
        with torch.no_grad():
            result = self.detect_oop(pressure)
        
        # 滑动窗口平滑
        self.history.append(result)
        if len(self.history) > 30:  # 1s @ 30Hz
            self.history.pop(0)
        
        # 连续3帧OOP才触发
        oop_count = sum(1 for r in self.history[-3:] if r['is_oop'])
        
        if oop_count >= 3:
            return {
                'alert': 'OOP_DETECTED',
                'level': 'LEVEL_2',
                'type': result['pressure_distribution'],
                'confidence': result['confidence']
            }
        elif oop_count >= 1:
            return {
                'alert': 'OOP_WARNING',
                'level': 'LEVEL_1',
                'type': result['pressure_distribution'],
                'confidence': result['confidence']
            }
        return {'alert': None, 'level': 'NORMAL', 'confidence': 1.0}