"""
智能座舱座椅人体姿态识别
基于 Wiley 2026 论文

核心: 2D姿态→3D重建→姿态分类→座椅参数优化
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Dict, List, Tuple, Optional

class PoseEstimator2D(nn.Module):
    """
    2D人体姿态估计器
    
    从单张图像估计17个关节点的2D坐标
    基于HRNet架构（高分辨率保持）
    """
    
    # COCO 17关键点
    KEYPOINT_NAMES = [
        'nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear',
        'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow',
        'left_wrist', 'right_wrist', 'left_hip', 'right_hip',
        'left_knee', 'right_knee', 'left_ankle', 'right_ankle'
    ]
    
    # 骨架连接
    SKELETON = [
        (5, 7), (7, 9), (6, 8), (8, 10),  # 手臂
        (5, 6),  # 肩
        (5, 11), (6, 12), (11, 12),  # 躯干
        (11, 13), (13, 15), (12, 14), (14, 16),  # 腿
        (0, 5), (0, 6),  # 头到肩
    ]
    
    def __init__(self, num_keypoints: int = 17):
        super().__init__()
        self.num_kp = num_keypoints
        
        # 轻量级HRNet
        self.stem = nn.Sequential(
            nn.Conv2d(3, 32, 3, 2, 1), nn.BatchNorm2d(32), nn.ReLU(),
            nn.Conv2d(32, 64, 3, 2, 1), nn.BatchNorm2d(64), nn.ReLU(),
        )
        
        # 多分辨率分支
        self.branch1 = self._make_branch(64, 32)
        self.branch2 = self._make_branch(64, 32, stride=2)
        self.branch3 = self._make_branch(64, 32, stride=4)
        
        # 融合
        self.fusion = nn.Conv2d(96, 64, 1)
        
        # 热力图输出
        self.heatmap_head = nn.Sequential(
            nn.Conv2d(64, 64, 3, 1, 1), nn.ReLU(),
            nn.Conv2d(64, num_keypoints, 1),
        )
        
    def _make_branch(self, in_ch, out_ch, stride=1):
        layers = []
        s = stride
        for _ in range(4):
            layers.extend([
                nn.Conv2d(in_ch, out_ch, 3, s, 1),
                nn.BatchNorm2d(out_ch), nn.ReLU(),
            ])
            in_ch = out_ch
            s = 1
        return nn.Sequential(*layers)
    
    def forward(self, image: torch.Tensor) -> Dict:
        """
        Args:
            image: (B, 3, H, W)
        Returns:
            heatmaps: (B, 17, H/4, W/4)
            keypoints: (B, 17, 2) 坐标
            confidence: (B, 17)
        """
        x = self.stem(image)
        
        b1 = self.branch1(x)
        b2 = self.branch2(x)
        b3 = self.branch3(x)
        
        # 上采样对齐
        b2_up = F.interpolate(b2, size=b1.shape[2:], mode='bilinear')
        b3_up = F.interpolate(b3, size=b1.shape[2:], mode='bilinear')
        
        fused = self.fusion(torch.cat([b1, b2_up, b3_up], dim=1))
        heatmaps = self.heatmap_head(fused)
        
        # 提取关键点坐标
        B, K, H, W = heatmaps.shape
        heatmaps_flat = heatmaps.view(B, K, -1)
        max_idx = heatmaps_flat.argmax(dim=-1)
        
        x_coord = (max_idx % W).float() / W
        y_coord = (max_idx // W).float() / H
        keypoints = torch.stack([x_coord, y_coord], dim=-1)
        confidence = heatmaps_flat.max(dim=-1).values
        
        return {
            'heatmaps': heatmaps,
            'keypoints': keypoints,
            'confidence': confidence,
        }


class Pose3DReconstructor(nn.Module):
    """
    3D姿态重建器
    
    从2D关节推断3D位置
    用于OOP检测和座椅参数优化
    """
    
    def __init__(self, num_joints: int = 17, hidden_dim: int = 256):
        super().__init__()
        # 2D→3D提升网络
        self.lift = nn.Sequential(
            nn.Linear(num_joints * 2, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, num_joints * 3),
        )
        
    def forward(self, keypoints_2d: torch.Tensor) -> torch.Tensor:
        """
        Args:
            keypoints_2d: (B, 17, 2)
        Returns:
            keypoints_3d: (B, 17, 3) 相对骨盆的3D坐标
        """
        B, K, _ = keypoints_2d.shape
        x = keypoints_2d.reshape(B, -1)
        pred_3d = self.lift(x).reshape(B, K, 3)
        
        # 以骨盆（髋关节中点）为原点
        pelvis = (pred_3d[:, 11] + pred_3d[:, 12]) / 2
        keypoints_3d = pred_3d - pelvis.unsqueeze(1)
        
        return keypoints_3d


class PostureClassifier(nn.Module):
    """
    姿态分类器
    
    将3D姿态分类为标准/异常姿态
    用于OOP检测
    """
    
    POSTURE_CLASSES = [
        'normal_driving',      # 正常驾驶
        'leaning_left',        # 左倾
        'leaning_right',       # 右倾
        'leaning_forward',     # 前倾
        'leaning_back',        # 后仰
        'reclined',            # 躺靠
        'reaching_dashboard',  # 摸仪表盘
        'picking_floor',       # 捡东西
        'sleeping',            # 睡眠
        'on_phone',            # 手机
        'twisting',            # 扭转
    ]
    
    def __init__(self, num_joints: int = 17, num_classes: int = 11):
        super().__init__()
        # 3D关节特征
        self.classifier = nn.Sequential(
            nn.Linear(num_joints * 3, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, num_classes),
        )
        
    def forward(self, keypoints_3d: torch.Tensor) -> Dict:
        """
        Args:
            keypoints_3d: (B, 17, 3)
        Returns:
            class_logits: (B, num_classes)
            posture: str
            confidence: float
        """
        B = keypoints_3d.shape[0]
        x = keypoints_3d.reshape(B, -1)
        logits = self.classifier(x)
        probs = F.softmax(logits, dim=-1)
        
        idx = probs.argmax(dim=-1)
        results = []
        for i in range(B):
            results.append({
                'posture': self.POSTURE_CLASSES[idx[i].item()],
                'confidence': probs[i, idx[i]].item(),
            })
        
        return {'results': results, 'logits': logits}


class SeatAdjustmentOptimizer(nn.Module):
    """
    座椅参数优化器
    
    根据姿态自动调节座椅参数
    
    控制参数:
    - seat_position: 前后(x)、高低(z) 0-1
    - backrest_angle: 靠背角度 0-1 (0=垂直, 1=平躺)
    - lumbar_support: 腰托高度/凸度 0-1
    - headrest_height: 头枕高度 0-1
    - side_bolster: 侧支撑 0-1
    """
    
    SEAT_PARAMS = {
        'seat_x': {'range': (0.0, 1.0), 'default': 0.5},
        'seat_z': {'range': (0.0, 1.0), 'default': 0.5},
        'backrest_angle': {'range': (0.0, 0.7), 'default': 0.25},  # 0-~40°
        'lumbar_height': {'range': (0.0, 1.0), 'default': 0.5},
        'lumbar_convexity': {'range': (0.0, 1.0), 'default': 0.4},
        'headrest_height': {'range': (0.0, 1.0), 'default': 0.5},
        'side_bolster': {'range': (0.0, 1.0), 'default': 0.3},
    }
    
    def __init__(self, num_joints: int = 17):
        super().__init__()
        num_params = len(self.SEAT_PARAMS)
        
        self.optimizer = nn.Sequential(
            nn.Linear(num_joints * 3, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, num_params),
            nn.Sigmoid(),  # 输出0-1
        )
        
        # 性别/身高差分参数
        self.body_scale = nn.Sequential(
            nn.Linear(num_joints * 3, 64),
            nn.ReLU(),
            nn.Linear(64, 3),  # 身高估计、坐高比、肩宽
        )
        
    def forward(self, keypoints_3d: torch.Tensor) -> Dict:
        """
        根据姿态生成座椅参数
        
        Args:
            keypoints_3d: (B, 17, 3)
            
        Returns:
            params: 座椅参数字典
            body_estimate: 身体尺寸估计
        """
        B = keypoints_3d.shape[0]
        x = keypoints_3d.reshape(B, -1)
        
        raw_params = self.optimizer(x)
        body = self.body_scale(x)
        
        # 映射到具体参数范围
        results = []
        param_names = list(self.SEAT_PARAMS.keys())
        
        for i in range(B):
            params = {}
            for j, name in enumerate(param_names):
                min_val = self.SEAT_PARAMS[name]['range'][0]
                max_val = self.SEAT_PARAMS[name]['range'][1]
                value = raw_params[i, j].item() * (max_val - min_val) + min_val
                params[name] = round(value, 3)
            results.append(params)
        
        body_estimates = []
        for i in range(B):
            body_estimates.append({
                'estimated_height_cm': 150 + body[i, 0].item() * 50,  # 150-200cm
                'sitting_ratio': body[i, 1].item(),
                'shoulder_width_cm': 35 + body[i, 2].item() * 15,  # 35-50cm
            })
        
        return {
            'seat_params': results,
            'body_estimate': body_estimates,
        }


# ========== 完整系统 ==========
class IntelligentSeatSystem(nn.Module):
    """
    智能座椅系统
    
    集成: 2D姿态估计 → 3D重建 → 姿态分类 → 座椅优化
    同时: OOP异常检测
    """
    
    def __init__(self):
        super().__init__()
        self.pose2d = PoseEstimator2D()
        self.pose3d = Pose3DReconstructor()
        self.posture_cls = PostureClassifier()
        self.seat_opt = SeatAdjustmentOptimizer()
        
    def forward(self, image: torch.Tensor) -> Dict:
        """
        完整推理流程
        
        Args:
            image: (B, 3, H, W) 座舱侧摄像头图像
            
        Returns:
            result: 姿态、姿态分类、座椅参数、OOP状态
        """
        # 1. 2D姿态
        pose2d_out = self.pose2d(image)
        
        # 2. 3D重建
        kpts_3d = self.pose3d(pose2d_out['keypoints'])
        
        # 3. 姿态分类
        posture_out = self.posture_cls(kpts_3d)
        
        # 4. 座椅优化
        seat_out = self.seat_opt(kpts_3d)
        
        # 5. OOP判断
        oop_results = []
        for r in posture_out['results']:
            is_oop = r['posture'] not in [
                'normal_driving', 'leaning_left', 'leaning_right'
            ]
            oop_results.append({
                'is_oop': is_oop,
                'posture': r['posture'],
                'confidence': r['confidence'],
            })
        
        return {
            'keypoints_2d': pose2d_out['keypoints'],
            'keypoints_3d': kpts_3d,
            'confidence': pose2d_out['confidence'],
            'posture': posture_out['results'],
            'seat_params': seat_out['seat_params'],
            'body_estimate': seat_out['body_estimate'],
            'oop_status': oop_results,
        }


# ========== 测试 ==========
if __name__ == "__main__":
    torch.manual_seed(42)
    
    system = IntelligentSeatSystem()
    
    # 模拟座舱图像
    image = torch.randn(2, 3, 256, 192)  # 侧摄像头
    
    result = system(image)
    
    for i in range(2):
        print(f"\n{'='*60}")
        print(f"乘员 {i+1}:")
        print(f"  2D关键点: {result['keypoints_2d'][i].shape}")
        print(f"  3D关键点: {result['keypoints_3d'][i].shape}")
        print(f"  平均置信度: {result['confidence'][i].mean():.3f}")
        print(f"  姿态: {result['posture'][i]['posture']} "
              f"({result['posture'][i]['confidence']:.3f})")
        print(f"  OOP异常: {result['oop_status'][i]['is_oop']}")
        
        params = result['seat_params'][i]
        print(f"  座椅参数:")
        for k, v in params.items():
            print(f"    {k}: {v}")
        
        body = result['body_estimate'][i]
        print(f"  身体估计: 身高{body['estimated_height_cm']:.0f}cm, "
              f"肩宽{body['shoulder_width_cm']:.0f}cm")
    
    print(f"\n✅ 测试通过")