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

class PostureResNet(nn.Module):
    """
    座椅压力分布坐姿识别ResNet
    
    论文: Applied Sciences, 2025
    输入: 42×48压力分布矩阵
    输出: 9类坐姿
    准确率: 94.2%（混合座面训练）
    
    关键: 压力分布→2D图像→CNN处理
    """
    
    def __init__(self, num_postures=9):
        super().__init__()
        
        # 输入：1×42×48（压力矩阵）
        self.conv1 = nn.Conv2d(1, 64, 7, stride=2, padding=3)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
        
        # ResNet块
        self.layer1 = self._make_layer(64, 64, 2, stride=1)
        self.layer2 = self._make_layer(64, 128, 2, stride=2)
        self.layer3 = self._make_layer(128, 256, 2, stride=2)
        
        # 分类头
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(256, num_postures)
        
    def _make_layer(self, in_ch, out_ch, blocks, stride):
        layers = []
        layers.append(self._basic_block(in_ch, out_ch, stride))
        for _ in range(1, blocks):
            layers.append(self._basic_block(out_ch, out_ch, 1))
        return nn.Sequential(*layers)
    
    def _basic_block(self, in_ch, out_ch, stride):
        return ResidualBlock(in_ch, out_ch, stride)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x


class ResidualBlock(nn.Module):
    def __init__(self, in_ch, out_ch, stride=1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_ch)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_ch)
        
        self.shortcut = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 1, stride=stride, bias=False),
            nn.BatchNorm2d(out_ch)
        ) if stride != 1 or in_ch != out_ch else nn.Identity()
        
    def forward(self, x):
        out = torch.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        return torch.relu(out)


class SeatPosturePipeline:
    """
    座椅压力坐姿检测完整管道
    
    输入: 压力传感阵列
    输出: 坐姿分类 + OOP风险评估
    
    应用: Euro NCAP OOP检测要求
    """
    
    POSTURE_RISK = {
        'P1_normal': {'risk': 'safe', 'airbag': 'normal'},
        'P2_forward': {'risk': 'medium', 'airbag': 'reduce_power'},
        'P3_backward': {'risk': 'medium', 'airbag': 'normal'},
        'P4_left': {'risk': 'medium', 'airbag': 'disable_side'},
        'P5_right': {'risk': 'medium', 'airbag': 'disable_side'},
        'P6_cross_left': {'risk': 'high', 'airbag': 'disable'},
        'P7_cross_right': {'risk': 'high', 'airbag': 'disable'},
        'P8_sideways': {'risk': 'high', 'airbag': 'disable'},
        'P9_on_dashboard': {'risk': 'critical', 'airbag': 'disable'},
    }
    
    def __init__(self):
        self.model = PostureResNet()
        self.model.eval()
        
    def classify_posture(self, pressure_matrix: np.ndarray) -> dict:
        """
        分类坐姿
        
        Args:
            pressure_matrix: 42×48 压力分布矩阵 (0-100 kPa)
            
        Returns:
            dict: 坐姿分类 + 风险评估
        """
        # 归一化到0-1
        normalized = pressure_matrix / 100.0
        x = torch.FloatTensor(normalized).unsqueeze(0).unsqueeze(0)  # (1, 1, 42, 48)
        
        with torch.no_grad():
            logits = self.model(x)
            probs = torch.softmax(logits, dim=1)
            pred_idx = probs.argmax(1).item()
            confidence = probs[0, pred_idx].item()
        
        posture_key = f"P{pred_idx+1}"
        posture_names = ['normal', 'forward', 'backward', 'left', 'right',
                        'cross_left', 'cross_right', 'sideways', 'on_dashboard']
        posture_name = posture_names[pred_idx]
        
        risk = self.POSTURE_RISK[f'P{pred_idx+1}_{posture_name}']
        
        return {
            'posture': posture_name,
            'posture_id': posture_key,
            'confidence': confidence,
            'risk_level': risk['risk'],
            'airbag_action': risk['airbag'],
            'all_probs': {f'P{i+1}_{posture_names[i]}': probs[0, i].item() 
                         for i in range(9)}
        }


# 实际测试
if __name__ == "__main__":
    pipeline = SeatPosturePipeline()
    
    # 模拟正常坐姿压力分布
    np.random.seed(42)
    pressure = np.zeros((42, 48))
    # 臀部区域（中间偏下）高压力
    pressure[25:35, 15:35] = np.random.uniform(40, 80, (10, 20))
    # 背部区域（上部）中等压力
    pressure[5:15, 15:35] = np.random.uniform(20, 40, (10, 20))
    # 大腿区域
    pressure[35:42, 10:40] = np.random.uniform(10, 30, (7, 30))
    
    result = pipeline.classify_posture(pressure)
    
    print("=== 座椅压力坐姿检测 ===")
    print(f"坐姿: {result['posture']} ({result['posture_id']})")
    print(f"置信度: {result['confidence']:.2%}")
    print(f"风险: {result['risk_level']}")
    print(f"安全气囊: {result['airbag_action']}")
    
    print(f"\n=== 论文性能 ===")
    print(f"| 架构 | 硬座准确率 | 软座准确率 | 混合训练 |")
    print(f"|------|-----------|-----------|---------|")
    print(f"| FNN  | 87.3%     | 84.1%     | 82.5%   |")
    print(f"| CNN  | 91.5%     | 88.7%     | 89.3%   |")
    print(f"| ResNet| 94.2%    | 92.8%     | 94.2%   |")
    print(f"\nResNet在混合训练下最优: 94.2%")