"""
Rajagopal 全身骨骼模型实现

论文 Section II-B: 20个刚体段，37个广义坐标
- 6自由度骨盆浮动关节
- 31个铰链关节（解剖学定义旋转轴）
"""

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

class BiomechanicalSkeleton(nn.Module):
    """
    全身骨骼模型（基于 Rajagopal 2015）
    
    拓扑结构：
        pelvis (root, 6-DOF)
        ├── spine → thorax → head
        ├── left_hip → left_thigh → left_shin → left_foot
        ├── right_hip → right_thigh → right_shin → right_foot
        ├── left_shoulder → left_upper_arm → left_forearm → left_hand
        └── right_shoulder → right_upper_arm → right_forearm → right_hand
    
    20个刚体段，37个广义坐标
    """
    
    # 关节定义（名称, 父段, 旋转轴, 轴数量）
    JOINT_DEFINITION = {
        'pelvis': ('root', None, 6),        # 6-DOF浮动关节
        'spine': ('pelvis', 'sagittal', 3),  # 3轴弯曲
        'thorax': ('spine', 'sagittal', 3),
        'head': ('thorax', 'sagittal', 3),
        'left_hip': ('pelvis', 'sagittal', 3),
        'left_knee': ('left_thigh', 'sagittal', 1),  # 1轴铰链
        'left_ankle': ('left_shin', 'sagittal', 2),
        'left_shoulder': ('thorax', 'sagittal', 3),
        'left_elbow': ('left_upper_arm', 'sagittal', 1),
        'left_wrist': ('left_forearm', 'sagittal', 2),
        'right_hip': ('pelvis', 'sagittal', 3),
        'right_knee': ('right_thigh', 'sagittal', 1),
        'right_ankle': ('right_shin', 'sagittal', 2),
        'right_shoulder': ('thorax', 'sagittal', 3),
        'right_elbow': ('right_upper_arm', 'sagittal', 1),
        'right_wrist': ('right_forearm', 'sagittal', 2),
    }
    
    # 段长比例（相对于身高的归一化值，来自 Rajagopal 2015）
    SEGMENT_RATIOS = {
        'pelvis_to_spine': 0.282,
        'spine_to_thorax': 0.221,
        'thorax_to_head': 0.182,
        'hip_to_knee': 0.246,
        'knee_to_ankle': 0.245,
        'shoulder_to_elbow': 0.186,
        'elbow_to_wrist': 0.146,
    }
    
    def __init__(self, num_markers: int = 17, num_foot_bodies: int = 4):
        super().__init__()
        self.num_q = 37  # 广义坐标数
        self.num_bodies = 20
        self.num_markers = num_markers
        self.num_foot_bodies = num_foot_bodies
        
        # 默认段长（将根据预测的缩放因子调整）
        self.register_buffer(
            'default_segment_lengths',
            torch.tensor(list(self.SEGMENT_RATIOS.values()))
        )
        
    def forward_kinematics(
        self,
        q: torch.Tensor,           # (B, T, 37) 广义坐标
        segment_scales: torch.Tensor  # (B, num_segments) 段长缩放
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        可微前向运动学
        
        Args:
            q: 广义坐标 (B, T, 37)
            segment_scales: 段长缩放因子
            
        Returns:
            markers: 3D关节位置 (B, T, 17, 3)
            foot_contacts: 接触概率 (B, T, 4)
        """
        B, T, _ = q.shape
        
        # 分解广义坐标
        pelvis_pos = q[..., :3]      # 骨盆位置
        pelvis_rot = q[..., 3:6]     # 骨盆旋转（欧拉角）
        joint_angles = q[..., 6:]   # 31个关节角度
        
        # 按骨骼拓扑顺序计算各段位置
        # （简化实现，实际需要按运动学树遍历）
        markers = self._compute_markers(
            pelvis_pos, pelvis_rot, joint_angles, segment_scales
        )
        
        # 脚-地接触分类
        foot_contacts = self._compute_foot_contact(markers)
        
        return markers, foot_contacts
    
    def _compute_markers(self, pos, rot, angles, scales):
        """
        简化的前向运动学计算
        
        实际实现中需要完整的运动学树遍历
        这里给出关键步骤的伪代码框架
        """
        # 这里用简化的线性组合代替完整 FK
        # 实际论文中使用了 PyTorch 实现的可微 FK
        batch_size, seq_len, _ = pos.shape
        markers = torch.zeros(batch_size, seq_len, self.num_markers, 3)
        
        # 骨盆 → 髋关节
        markers[:, :, 0, :] = pos  # 骨盆中心
        
        # 脊柱
        spine_len = scales[:, 0:1] * self.default_segment_lengths[0]
        markers[:, :, 1, :] = pos + torch.cat([
            torch.zeros_like(spine_len), 
            torch.zeros_like(spine_len), 
            spine_len
        ], dim=-1).unsqueeze(1)
        
        # 后续关节按运动学链顺序计算...
        # （完整实现需要 20 段的完整遍历）
        
        return markers
    
    def _compute_foot_contact(self, markers):
        """
        脚-地接触分类
        
        4个脚体：左/右跟骨、左/右脚趾
        判断依据：脚高度 < 阈值 且 脚速度 < 阈值
        """
        foot_indices = [7, 8, 10, 11]  # 左/右跟骨、脚趾索引
        foot_pos = markers[:, :, foot_indices, :]  # (B, T, 4, 3)
        
        # 脚高度
        foot_height = foot_pos[..., 2]  # z坐标
        
        # 脚速度
        foot_vel = torch.diff(
            foot_pos, dim=1, 
            prepend=foot_pos[:, :1, :, :]
        )
        foot_speed = torch.norm(foot_vel, dim=-1)
        
        # 接触概率（简化版）
        height_score = torch.sigmoid(-10 * (foot_height - 0.05))
        speed_score = torch.sigmoid(-10 * (foot_speed - 0.1))
        contact_prob = height_score * speed_score
        
        return contact_prob


class ScalePredictor(nn.Module):
    """
    从雷达点云特征预测主体特异性段长缩放
    
    论文 Section II-D: 主体特异性身体比例预测
    """
    
    def __init__(self, point_dim: int = 7, num_segments: int = 7):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(point_dim, 128),
            nn.GELU(),
            nn.Linear(128, 256),
            nn.GELU(),
            nn.Linear(256, 128),
            nn.GELU()
        )
        # 全局池化后预测段长
        self.scale_head = nn.Sequential(
            nn.Linear(128, 64),
            nn.GELU(),
            nn.Linear(64, num_segments),
            nn.Sigmoid()  # 归一化到 [0, 1]
        )
        
    def forward(
        self, 
        point_cloud: torch.Tensor,  # (B, N, F)
        mask: torch.Tensor = None    # (B, N)
    ) -> torch.Tensor:
        """
        Returns:
            scales: (B, num_segments) 段长缩放因子
        """
        x = self.encoder(point_cloud)  # (B, N, 128)
        
        if mask is not None:
            x = x * mask.unsqueeze(-1).float()
            x = x.sum(dim=1) / mask.sum(dim=1, keepdim=True).float()
        else:
            x = x.mean(dim=1)
            
        scales = self.scale_head(x)  # (B, num_segments)
        return scales


class MotionPredictionNetwork(nn.Module):
    """
    时序雷达 → 广义坐标预测网络
    
    论文 Section II-E: BiLSTM + 潜在递归
    """
    
    def __init__(
        self,
        point_dim: int = 7,
        hidden_dim: int = 256,
        num_q: int = 37,
        num_contacts: int = 4,
        num_layers: int = 2
    ):
        super().__init__()
        # 点云特征提取
        self.point_encoder = nn.Sequential(
            nn.Linear(point_dim, 128),
            nn.GELU(),
            nn.Linear(128, hidden_dim)
        )
        
        # 空间聚合
        self.spatial_attn = nn.MultiheadAttention(
            hidden_dim, num_heads=8, batch_first=True
        )
        
        # 时序 BiLSTM
        self.temporal_lstm = nn.LSTM(
            input_size=hidden_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=True
        )
        
        # 输出头
        self.q_head = nn.Linear(hidden_dim * 2, num_q)  # 广义坐标
        self.contact_head = nn.Linear(hidden_dim * 2, num_contacts)  # 接触logits
        
    def forward(
        self,
        radar_sequence: torch.Tensor,  # (B, T, N, F)
        point_mask: torch.Tensor = None  # (B, T, N)
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Returns:
            q: (B, T, 37) 广义坐标
            contact_logits: (B, T, 4) 接触分类logits
        """
        B, T, N, F = radar_sequence.shape
        
        # 逐帧编码
        x = self.point_encoder(radar_sequence)  # (B, T, N, H)
        x = x.view(B * T, N, -1)
        
        if point_mask is not None:
            point_mask_flat = point_mask.view(B * T, N)
        else:
            point_mask_flat = None
            
        x, _ = self.spatial_attn(x, x, x, key_padding_mask=point_mask_flat)
        x = x.mean(dim=1)  # 帧级聚合
        x = x.view(B, T, -1)
        
        # 时序建模
        x, _ = self.temporal_lstm(x)
        
        # 输出
        q = self.q_head(x)  # (B, T, 37)
        contact_logits = self.contact_head(x)  # (B, T, 4)
        
        return q, contact_logits


class RadarBiomechanicalHPE(nn.Module):
    """
    完整的雷达-生物力学人体姿态估计框架
    
    端到端可微管道：
        雷达点云 → MotionPrediction → 广义坐标
        → BiomechanicalSkeleton.forward_kinematics → 3D关节位置 + 接触
        → 联合损失（位置 + 接触）
    """
    
    def __init__(self):
        super().__init__()
        self.motion_net = MotionPredictionNetwork()
        self.scale_net = ScalePredictor()
        self.skeleton = BiomechanicalSkeleton()
        
    def forward(
        self,
        radar_seq: torch.Tensor,      # (B, T, N, F)
        point_mask: torch.Tensor = None  # (B, T, N)
    ) -> Dict[str, torch.Tensor]:
        # 预测广义坐标和接触
        q, contact_logits = self.motion_net(radar_seq, point_mask)
        
        # 从第一帧预测段长缩放
        first_frame = radar_seq[:, 0, :, :]  # (B, N, F)
        first_mask = point_mask[:, 0, :] if point_mask is not None else None
        scales = self.scale_net(first_frame, first_mask)
        
        # 前向运动学
        markers, foot_contacts = self.skeleton.forward_kinematics(q, scales)
        
        return {
            'q': q,                          # 广义坐标
            'contact_logits': contact_logits,  # 接触logits
            'scales': scales,                 # 段长缩放
            'markers': markers,               # 3D关节位置
            'foot_contacts': foot_contacts     # 脚接触概率
        }


class CombinedLoss(nn.Module):
    """
    组合损失：数据追踪 + 接触分类
    
    L = L_pos + λ_contact * L_contact + λ_scale * L_scale
    """
    
    def __init__(self, lambda_contact=1.0, lambda_scale=0.5):
        super().__init__()
        self.lambda_contact = lambda_contact
        self.lambda_scale = lambda_scale
        
    def forward(
        self,
        pred: dict,
        gt_markers: torch.Tensor,      # (B, T, 17, 3)
        gt_contacts: torch.Tensor,     # (B, T, 4)
        gt_scales: torch.Tensor        # (B, 7)
    ) -> Tuple[torch.Tensor, dict]:
        # 位置损失
        pos_loss = nn.functional.mse_loss(pred['markers'], gt_markers)
        
        # 接触分类损失
        contact_loss = nn.functional.binary_cross_entropy_with_logits(
            pred['contact_logits'], gt_contacts.float()
        )
        
        # 段长缩放损失
        scale_loss = nn.functional.mse_loss(pred['scales'], gt_scales)
        
        total = (
            pos_loss 
            + self.lambda_contact * contact_loss
            + self.lambda_scale * scale_loss
        )
        
        return total, {
            'pos_loss': pos_loss.item(),
            'contact_loss': contact_loss.item(),
            'scale_loss': scale_loss.item()
        }


# ============ 实验复现 ============
if __name__ == "__main__":
    torch.manual_seed(42)
    
    # 模拟输入：11个受试者康复训练数据
    B, T, N, F = 4, 30, 192, 7  # 4个样本, 30帧, 192点, 7特征
    
    radar_input = torch.randn(B, T, N, F)
    point_mask = torch.ones(B, T, N, dtype=torch.bool)
    
    # GT
    gt_markers = torch.randn(B, T, 17, 3) * 1.0  # 米
    gt_contacts = torch.randint(0, 2, (B, T, 4)).float()
    gt_scales = torch.rand(B, 7) * 0.3 + 0.85  # [0.85, 1.15]
    
    # 初始化模型
    model = RadarBiomechanicalHPE()
    criterion = CombinedLoss()
    
    # 前向
    pred = model(radar_input, point_mask)
    
    # 损失
    loss, metrics = criterion(pred, gt_markers, gt_contacts, gt_scales)
    
    print("=" * 60)
    print("Radar-Biomechanical HPE 测试结果")
    print("=" * 60)
    print(f"输入: 雷达点云 {radar_input.shape}")
    print(f"广义坐标: {pred['q'].shape}")
    print(f"关节位置: {pred['markers'].shape}")
    print(f"接触logits: {pred['contact_logits'].shape}")
    print(f"段长缩放: {pred['scales'].shape}")
    print(f"\n损失: {loss.item():.4f}")
    for k, v in metrics.items():
        print(f"  {k}: {v:.4f}")
    
    # 参数量
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n模型参数: {total_params:,} ({total_params/1e6:.2f}M)")

# 车内 OOP 检测方案
oop_config = {
    "sensor": "TI IWR6843 60GHz",
    "mount": "车顶内衬",
    "frame_rate": "15fps",
    "points_per_frame": "≤64",
    "superframe": "3帧聚合 → 192点",
    
    # OOP 分类
    "normal_postures": [
        "正常坐姿 (q接近零位)",
        "前倾 (躯干前屈角度 > 15°)",
        "后仰 (躯干后仰 > 10°)",
        "侧倾 (躯干侧弯 > 10°)"
    ],
    "abnormal_postures": [
        "跪姿 (膝关节屈曲 > 120°, 骨盆下降)",
        "趴姿 (躯干前屈 > 45°, 头部低于肩)",
        "站立 (骨盆z > 座椅高度)",
        "蜷缩 (多关节同时屈曲)"
    ],
    
    # 判定逻辑
    "detection_logic": """
        1. 从广义坐标 q 提取关键角度
        2. 躯干前屈角 = q[spine_flexion]
        3. 髋关节角 = q[hip_flexion_left/right]
        4. 膝关节角 = q[knee_flexion_left/right]
        5. 骨盆高度 = q[pelvis_z]
        6. 与正常坐姿范围比较 → 偏差 > 阈值 → 异常
    """,
    
    "advantages": [
        "关节角度直接可解释（不需要后处理）",
        "骨骼长度固定（不会出现解剖学不可能姿态）",
        "脚-地接触判断（区分站立vs坐姿）",
        "隐私保护（无图像）"
    ]
}