"""
智能座椅压力分布姿态识别
基于深度学习（改进 LeNet-5 + 薄膜压力阵列）

参考论文：
1. Wang B et al. (2025) "Deep Learning-Based Sitting Posture Recognition 
   from Pressure Distribution Across Hard and Soft Seat Environment"
   Applied Sciences 15(23):12744
2. MDPI Processes (2026) "Human Sitting Posture Recognition Based on 
   Improved LeNet-5 Optimization Algorithm"

关键指标：
- 8种坐姿分类
- 96.3% 准确率（改进LeNet-5）
- 推理延迟 < 10ms
- 传感器：薄膜压力阵列 44×52 = 2288 传感单元
"""

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

class ImprovedLeNet5(nn.Module):
    """
    改进 LeNet-5 用于坐姿识别
    
    改进点（参考 Processes 2026 论文）：
    1. 输入从 32×32 扩展到 44×52（压力垫尺寸）
    2. 增加 BatchNorm 层提升训练稳定性
    3. ReLU 替代原始 Sigmoid
    4. Dropout 防止过拟合
    5. 输出从 10 类改为 8 种坐姿
    
    坐姿类别：
    0: 正常坐姿（理想姿态）
    1: 前倾（看手机/操作）
    2: 后仰（放松/疲劳）
    3: 左倾（侧靠）
    4: 右倾（侧靠）
    5: 跷二郎腿（左）
    6: 跷二郎腿（右）
    7: 滑动/潜艇姿态（OOP风险）
    """
    
    def __init__(self, num_classes: int = 8, input_channels: int = 1):
        super().__init__()
        
        # 卷积层1：提取局部压力模式
        self.conv1 = nn.Conv2d(input_channels, 6, kernel_size=5, padding=2)
        self.bn1 = nn.BatchNorm2d(6)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # 卷积层2：提取更高层特征
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5, padding=2)
        self.bn2 = nn.BatchNorm2d(16)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # 卷积层3（新增）：增强特征提取
        self.conv3 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(32)
        self.relu3 = nn.ReLU()
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # 全连接层
        self.fc1 = nn.Linear(32 * 5 * 6, 120)  # 44×52 → 5×6 after 3 pools
        self.relu4 = nn.ReLU()
        self.dropout1 = nn.Dropout(0.3)
        
        self.fc2 = nn.Linear(120, 84)
        self.relu5 = nn.ReLU()
        self.dropout2 = nn.Dropout(0.3)
        
        self.fc3 = nn.Linear(84, num_classes)
    
    def forward(self, x):
        # x: (B, 1, 44, 52)
        x = self.pool1(self.relu1(self.bn1(self.conv1(x))))
        x = self.pool2(self.relu2(self.bn2(self.conv2(x))))
        x = self.pool3(self.relu3(self.bn3(self.conv3(x))))
        
        x = x.view(x.size(0), -1)
        x = self.dropout1(self.relu4(self.fc1(x)))
        x = self.dropout2(self.relu5(self.fc2(x)))
        x = self.fc3(x)
        
        return x


class SmartSeatSystem:
    """
    智能座椅系统集成层
    
    融合压力分布、生物计量、温度、位置传感器
    实现舒适+安全双目标
    """
    
    # 8种坐姿的安全/舒适评估
    POSTURE_PROFILES = {
        0: {'name': '正常坐姿', 'comfort': 'optimal', 'safety': 'safe', 
            'oop_risk': 'none', 'action': 'maintain'},
        1: {'name': '前倾', 'comfort': 'poor', 'safety': 'moderate',
            'oop_risk': 'moderate', 'action': 'lumbar_increase'},
        2: {'name': '后仰', 'comfort': 'fair', 'safety': 'moderate',
            'oop_risk': 'low', 'action': 'backrest_forward'},
        3: {'name': '左倾', 'comfort': 'poor', 'safety': 'risk',
            'oop_risk': 'moderate', 'action': 'bolster_left'},
        4: {'name': '右倾', 'comfort': 'poor', 'safety': 'risk',
            'oop_risk': 'moderate', 'action': 'bolster_right'},
        5: {'name': '跷二郎腿(左)', 'comfort': 'fair', 'safety': 'risk',
            'oop_risk': 'high', 'action': 'alert_reposition'},
        6: {'name': '跷二郎腿(右)', 'comfort': 'fair', 'safety': 'risk',
            'oop_risk': 'high', 'action': 'alert_reposition'},
        7: {'name': '潜艇姿态', 'comfort': 'poor', 'safety': 'danger',
            'oop_risk': 'critical', 'action': 'critical_alert'},
    }
    
    def __init__(self, config: dict):
        # 压力垫配置
        self.pressure_rows = config.get('rows', 44)
        self.pressure_cols = config.get('cols', 52)
        self.pressure_fps = config.get('fps', 10)
        
        # 姿态分类模型
        self.posture_model = ImprovedLeNet5(num_classes=8)
        self.posture_model.eval()
        
        # 生理参数
        self.heart_rate = 0
        self.breathing_rate = 0
        self.hrv = 0  # 心率变异性
        
        # 舒适度基线
        self.comfort_baseline = None
        self.session_start_time = None
        
        # 疲劳指标
        self.fatigue_score = 0.0
        self.posture_change_count = 0
        self.time_in_poor_posture = 0
        
    def process_pressure_frame(self, pressure_data: np.ndarray) -> dict:
        """
        处理单帧压力数据
        
        Args:
            pressure_data: (44, 52) 压力分布矩阵 (kPa)
        
        Returns:
            analysis: 姿态分类 + COF + 对称性 + 舒适度
        """
        # 归一化到 [0, 1]
        normalized = pressure_data / 100.0  # 假设最大100 kPa
        normalized = np.clip(normalized, 0, 1)
        
        # 转为模型输入
        model_input = torch.FloatTensor(normalized).unsqueeze(0).unsqueeze(0)
        
        # 姿态分类
        with torch.no_grad():
            logits = self.posture_model(model_input)
            probabilities = torch.softmax(logits, dim=1)
            posture_class = torch.argmax(probabilities, dim=1).item()
            confidence = probabilities[0, posture_class].item()
        
        # 压力中心（COF）
        total_pressure = np.sum(pressure_data)
        if total_pressure > 0:
            x_coords = np.arange(self.pressure_cols) * 0.0116  # 11.6mm spacing
            y_coords = np.arange(self.pressure_rows) * 0.0116
            X, Y = np.meshgrid(x_coords, y_coords)
            
            cof_x = np.sum(X * pressure_data) / total_pressure
            cof_y = np.sum(Y * pressure_data) / total_pressure
        else:
            cof_x, cof_y = 0, 0
        
        # 左右对称性
        left_pressure = np.sum(pressure_data[:, :self.pressure_cols//2])
        right_pressure = np.sum(pressure_data[:, self.pressure_cols//2:])
        symmetry_ratio = min(left_pressure, right_pressure) / max(left_pressure, right_pressure + 1e-6)
        
        # 前后分布
        front_pressure = np.sum(pressure_data[:self.pressure_rows//2, :])
        back_pressure = np.sum(pressure_data[self.pressure_rows//2:, :])
        front_back_ratio = front_pressure / (back_pressure + 1e-6)
        
        # 舒适度评估
        posture_profile = self.POSTURE_PROFILES[posture_class]
        
        return {
            'posture_class': posture_class,
            'posture_name': posture_profile['name'],
            'confidence': confidence,
            'cof': (cof_x, cof_y),
            'symmetry_ratio': symmetry_ratio,
            'front_back_ratio': front_back_ratio,
            'total_pressure': total_pressure,
            'comfort_level': posture_profile['comfort'],
            'safety_level': posture_profile['safety'],
            'oop_risk': posture_profile['oop_risk'],
            'recommended_action': posture_profile['action'],
        }
    
    def monitor_fatigue(self, posture_history: List[dict],
                         duration_min: int = 30) -> dict:
        """
        通过座椅传感器监控疲劳
        
        疲劳指标：
        1. 姿态变化频率降低（疲劳时少动）
        2. 前倾比例增加（躯干前倒）
        3. 压力分布集中化（肌肉松弛）
        4. 心率下降趋势
        5. 呼吸率变化
        
        Args:
            posture_history: 历史姿态记录
            duration_min: 分析窗口（分钟）
        
        Returns:
            fatigue_assessment
        """
        if len(posture_history) < 10:
            return {'fatigue_score': 0, 'confidence': 0}
        
        # 1. 姿态变化频率
        posture_changes = 0
        for i in range(1, len(posture_history)):
            if posture_history[i]['posture_class'] != posture_history[i-1]['posture_class']:
                posture_changes += 1
        
        expected_changes = duration_min * 2  # 正常每分钟 ~2 次微调
        change_ratio = posture_changes / max(expected_changes, 1)
        # 疲劳时变化频率降低
        movement_score = max(0, 1 - change_ratio)  # 0=正常, 1=无移动
        
        # 2. 前倾比例
        forward_count = sum(
            1 for p in posture_history 
            if p['posture_class'] in [1, 7]  # 前倾/潜艇
        )
        forward_ratio = forward_count / len(posture_history)
        
        # 3. 对称性下降
        symmetry_trend = [p['symmetry_ratio'] for p in posture_history]
        symmetry_decline = max(0, symmetry_trend[0] - np.mean(symmetry_trend[-10:]))
        
        # 4. 综合疲劳得分
        fatigue_score = (
            0.35 * movement_score +
            0.30 * forward_ratio +
            0.20 * symmetry_decline +
            0.15 * max(0, 1 - self.heart_rate / 70)  # 心率偏低
        )
        
        fatigue_level = 'normal'
        if fatigue_score > 0.6:
            fatigue_level = 'high'
        elif fatigue_score > 0.3:
            fatigue_level = 'moderate'
        
        return {
            'fatigue_score': fatigue_score,
            'fatigue_level': fatigue_level,
            'movement_score': movement_score,
            'forward_ratio': forward_ratio,
            'posture_changes': posture_changes,
            'recommended_action': 'alert' if fatigue_score > 0.6 else 'monitor',
        }


# ==================== 测试 ====================
if __name__ == "__main__":
    np.random.seed(42)
    
    config = {'rows': 44, 'cols': 52, 'fps': 10}
    system = SmartSeatSystem(config)
    
    print(f"=== 智能座椅系统 ===")
    print(f"压力垫: {config['rows']}×{config['cols']} = {config['rows']*config['cols']} 传感器")
    print(f"采样率: {config['fps']} Hz")
    
    # 模拟正常坐姿压力分布
    normal_pressure = np.zeros((44, 52))
    # 两个大腿区域 + 臀部区域
    normal_pressure[25:40, 10:20] = 60  # 左大腿
    normal_pressure[25:40, 32:42] = 60  # 右大腿
    normal_pressure[35:44, 15:37] = 80  # 臀部
    normal_pressure += np.random.normal(0, 3, (44, 52))
    normal_pressure = np.clip(normal_pressure, 0, 100)
    
    # 模拟前倾（看手机）
    forward_pressure = np.zeros((44, 52))
    forward_pressure[15:30, 15:37] = 50  # 前移的大腿
    forward_pressure[5:20, 18:34] = 40  # 前倾的背
    forward_pressure += np.random.normal(0, 3, (44, 52))
    forward_pressure = np.clip(forward_pressure, 0, 100)
    
    # 模拟跷二郎腿
    cross_pressure = np.zeros((44, 52))
    cross_pressure[25:40, 10:25] = 70  # 重叠腿
    cross_pressure[25:40, 25:42] = 30  # 下方腿（轻）
    cross_pressure[35:44, 15:37] = 60
    cross_pressure += np.random.normal(0, 3, (44, 52))
    cross_pressure = np.clip(cross_pressure, 0, 100)
    
    # 分析各姿态
    for name, pressure in [('正常', normal_pressure), 
                           ('前倾', forward_pressure),
                           ('跷二郎腿', cross_pressure)]:
        result = system.process_pressure_frame(pressure)
        print(f"\n--- {name} ---")
        print(f"  分类: {result['posture_name']} (class={result['posture_class']})")
        print(f"  置信度: {result['confidence']:.3f}")
        print(f"  COF: ({result['cof'][0]*100:.1f}, {result['cof'][1]*100:.1f}) cm")
        print(f"  对称性: {result['symmetry_ratio']:.3f}")
        print(f"  舒适度: {result['comfort_level']}")
        print(f"  安全: {result['safety_level']}")
        print(f"  OOP风险: {result['oop_risk']}")
        print(f"  建议: {result['recommended_action']}")
    
    # 疲劳监控模拟
    history = []
    for i in range(100):
        if i < 30:
            history.append(system.process_pressure_frame(normal_pressure))
        elif i < 60:
            # 疲劳开始：前倾增加，变化减少
            history.append(system.process_pressure_frame(forward_pressure))
        else:
            history.append(system.process_pressure_frame(forward_pressure))
    
    fatigue = system.monitor_fatigue(history, duration_min=10)
    print(f"\n=== 疲劳监控 ===")
    print(f"疲劳得分: {fatigue['fatigue_score']:.3f}")
    print(f"疲劳等级: {fatigue['fatigue_level']}")
    print(f"运动得分: {fatigue['movement_score']:.3f}")
    print(f"前倾比例: {fatigue['forward_ratio']:.3f}")
    print(f"姿态变化: {fatigue['posture_changes']}")