# M4Human 数据采集硬件配置
hardware_config = {
    "radar": {
        "model": "Vayyar VSWR-10",  # 高分辨率 mmWave 雷达
        "band": "62-69 GHz",        # 7 GHz 带宽
        "channels": "高分辨率阵列",
        "output": ["raw_tensor", "point_cloud"]
    },
    "rgb": {
        "model": "Standard RGB Camera",
        "resolution": "1920x1080",
        "fps": 30
    },
    "depth": {
        "model": "Azure Kinect / RealSense",
        "resolution": "1280x1024",
        "fps": 30
    },
    "mocap": {
        "system": "Marker-based OptiTrack",
        "markers": 53,  # 全身标记点
        "mesh": "SMPL-X 3D mesh",
        "trajectory": "global 3D trajectory"
    }
}

import numpy as np
from typing import Tuple, Dict

class M4HumanRadarPipeline:
    """
    M4Human 雷达信号处理管道
    
    支持双模态输出：
    - Raw Tensor (RT): 保留完整信号信息，用于深度学习
    - Point Cloud (RPC): 经过处理后的点云，用于传统方法
    """
    
    def __init__(self, config: dict):
        self.bandwidth = config.get('bandwidth', 7e9)  # 7 GHz
        self.carrier_freq = config.get('carrier_freq', 65e9)  # 中心频率 65 GHz
        self.num_tx = config.get('num_tx', 16)
        self.num_rx = config.get('num_rx', 16)
        
    def compute_resolution(self) -> dict:
        """
        计算 M4Human 雷达空间分辨率
        
        Returns:
            resolutions: 距离/角度/速度分辨率
        """
        c = 3e8  # 光速
        
        # 距离分辨率
        range_res = c / (2 * self.bandwidth)
        
        # 角度分辨率（假设孔径 D = 0.1m）
        aperture = 0.1  # 10cm
        angular_res = 0.5 * 0.7 * (3e8 / self.carrier_freq) / aperture
        
        # 速度分辨率（假设观测时间 T = 50ms）
        obs_time = 0.05
        velocity_res = 0.5 * 3e8 / (2 * self.carrier_freq * obs_time)
        
        return {
            "range_resolution_m": round(range_res * 100, 2),    # cm
            "angular_resolution_deg": round(np.degrees(angular_res), 2),
            "velocity_resolution_cm_s": round(velocity_res * 100, 2),
            "bandwidth_GHz": self.bandwidth / 1e9
        }
    
    def raw_tensor_extraction(self, adc_data: np.ndarray) -> np.ndarray:
        """
        从 ADC 原始数据提取 Raw Tensor
        
        Args:
            adc_data: ADC 采样数据, shape=(num_rx, num_tx, num_chirps, num_samples)
            
        Returns:
            rt: Raw Tensor, shape=(range_bins, doppler_bins, angle_bins)
        """
        # Range-Doppler-Angle 三维张量
        # 1. Range FFT (沿采样轴)
        range_profile = np.fft.fft(adc_data, axis=-1)
        
        # 2. Doppler FFT (沿 chirp 轴)
        range_doppler = np.fft.fftshift(
            np.fft.fft(range_profile, axis=-2), axes=-2
        )
        
        # 3. Angle FFT (沿虚拟天线轴)
        range_doppler_angle = np.fft.fftshift(
            np.fft.fft(range_doppler, axis=0), axes=0
        )
        
        return range_doppler_angle
    
    def point_cloud_extraction(self, rt: np.ndarray) -> np.ndarray:
        """
        从 Raw Tensor 提取 Point Cloud
        
        Args:
            rt: Raw Tensor, shape=(angle_bins, doppler_bins, range_bins)
            
        Returns:
            rpc: Point Cloud, shape=(N, 5) = [x, y, z, doppler, intensity]
        """
        # CFAR 检测
        threshold = np.mean(rt) + 3 * np.std(rt)
        detections = np.argwhere(rt > threshold)
        
        # 转换为 3D 坐标
        points = []
        for det in detections:
            angle_bin, doppler_bin, range_bin = det
            range_val = range_bin * self.compute_resolution()["range_resolution_m"] / 100
            angle_val = (angle_bin - rt.shape[0]//2) * self.compute_resolution()["angular_resolution_deg"]
            
            x = range_val * np.cos(np.radians(angle_val))
            y = range_val * np.sin(np.radians(angle_val))
            z = 0  # 2D 雷达，z 从角度估计
            doppler = (doppler_bin - rt.shape[1]//2) * self.compute_resolution()["velocity_resolution_cm_s"] / 100
            
            points.append([x, y, z, doppler, rt[det]])
        
        return np.array(points)


# 实际测试
if __name__ == "__main__":
    pipeline = M4HumanRadarPipeline({})
    res = pipeline.compute_resolution()
    print("=== M4Human 雷达分辨率 ===")
    for k, v in res.items():
        print(f"  {k}: {v}")
    # 输出示例：
    # range_resolution_m: 2.14 cm
    # angular_resolution_deg: ~1.8°
    # velocity_resolution_cm_s: ~2.3 cm/s
    # bandwidth_GHz: 7.0

import torch
import torch.nn as nn

class RadarHMRBaseline(nn.Module):
    """
    mmWave 雷达人体网格重建基线模型
    
    基于 M4Human 论文中的 benchmark 设计，
    支持 RT（Raw Tensor）和 RPC（Point Cloud）两种输入。
    
    论文发现：
    - RT 模态优于 RPC（保留更多信息）
    - 多模态融合（RT+RGB-D）进一步提升性能
    - 快速运动和非约束动作仍是挑战
    """
    
    def __init__(self, modality: str = "rt", backbone: str = "resnet50"):
        super().__init__()
        self.modality = modality
        
        if modality == "rt":
            # Raw Tensor: 3D CNN 处理 Range-Doppler-Angle
            self.encoder = nn.Sequential(
                nn.Conv3d(1, 32, kernel_size=3, padding=1),
                nn.BatchNorm3d(32),
                nn.ReLU(),
                nn.MaxPool3d(2),
                
                nn.Conv3d(32, 64, kernel_size=3, padding=1),
                nn.BatchNorm3d(64),
                nn.ReLU(),
                nn.MaxPool3d(2),
                
                nn.Conv3d(64, 128, kernel_size=3, padding=1),
                nn.BatchNorm3d(128),
                nn.ReLU(),
                nn.AdaptiveAvgPool3d(1)
            )
            feat_dim = 128
            
        elif modality == "rpc":
            # Point Cloud: PointNet++ 处理
            self.encoder = PointNetPP(
                in_channels=5,  # x, y, z, doppler, intensity
                out_channels=128
            )
            feat_dim = 128
            
        # SMPL-X 参数回归头
        self.pose_head = nn.Linear(feat_dim, 156)   # 52 joints * 3
        self.shape_head = nn.Linear(feat_dim, 10)   # SMPL betas
        self.trans_head = nn.Linear(feat_dim, 3)    # 全局平移
        
    def forward(self, x: torch.Tensor) -> dict:
        """
        Args:
            x: 雷达输入
            - RT: (B, 1, A, D, R) = (B, 1, angle_bins, doppler_bins, range_bins)
            - RPC: (B, N, 5) = (B, num_points, 5)
        
        Returns:
            output: {
                pose: (B, 156)  关节角度,
                shape: (B, 10)  体型参数,
                trans: (B, 3)   全局平移
            }
        """
        feat = self.encoder(x)
        if self.modality == "rt":
            feat = feat.flatten(1)
        
        return {
            "pose": self.pose_head(feat),
            "shape": self.shape_head(feat),
            "trans": self.trans_head(feat)
        }


class PointNetPP(nn.Module):
    """简化版 PointNet++ 用于 RPC 处理"""
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(in_channels, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, out_channels)
        )
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (B, N, C)
        x = self.mlp(x)           # (B, N, out)
        x = x.max(dim=1)[0]      # (B, out) - max pooling
        return x

import torch
import torch.nn as nn
import numpy as np

class OOPDetectionFromRadar(nn.Module):
    """
    基于 mmWave 雷达 HMR 的 OOP 异常姿态检测器
    
    设计思路：
    1. 使用 M4Human 预训练 HMR 模型提取 3D 人体姿态
    2. 结合座椅参考几何，判断姿态是否异常
    3. 输出 OOP 类别和置信度
    
    硬件配置：
    | 组件 | 型号 | 参数 |
    |------|------|------|
    | mmWave 雷达 | Vayyar VSWR-10 或 TI IWR6843AOP | 60-64 GHz, 4Tx/4Rx |
    | 处理器 | Qualcomm QCS8255 | Hexagon NPU, 26 TOPS |
    | 更新频率 | 30 fps | 端到端推理 |
    """
    
    # OOP 类别定义（对标 Euro NCAP OOP 测试场景）
    OOP_CLASSES = {
        0: "normal_posture",        # 正常坐姿
        1: "forward_lean",          # 前倾
        2: "sideways_lean",         # 侧倾
        3: "fetal_position",        # 蜷缩/胎儿姿势
        4: "feet_on_dashboard",    # 脚放仪表板
        5: "kneeling",             # 跪姿
        6: "lying_down",            # 躺倒
        7: "backward_facing",      # 反向面朝后
    }
    
    def __init__(self, hmr_checkpoint: str = None, num_classes: int = 8):
        super().__init__()
        
        # 冻结的 HMR 特征提取器
        self.hmr = RadarHMRBaseline(modality="rt")
        if hmr_checkpoint:
            self.hmr.load_state_dict(torch.load(hmr_checkpoint))
        for param in self.hmr.parameters():
            param.requires_grad = False
            
        # OOP 分类头
        self.classifier = nn.Sequential(
            nn.Linear(156 + 10 + 3, 256),  # pose + shape + trans
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, num_classes)
        )
        
        # 座椅参考几何（标准乘员位置）
        self.register_buffer(
            'seat_reference',
            torch.tensor([
                [0.0, 0.0, 0.5],   # 头部参考
                [0.0, 0.0, 0.2],   # 躺椅参考
                [0.0, -0.3, 0.1],  # 腰部参考
                [0.0, -0.5, 0.0],  # 臀部参考
                [0.0, -0.7, -0.1], # 膝盖参考
                [0.0, -0.9, -0.2], # 脚踝参考
            ])
        )
    
    def forward(self, radar_input: torch.Tensor) -> dict:
        """
        Args:
            radar_input: (B, 1, A, D, R) Raw Tensor
        
        Returns:
            output: {
                pose: (B, 156)  3D 关节位置,
                shape: (B, 10)  体型参数,
                trans: (B, 3)   全局平移,
                oop_logits: (B, 8)  OOP 分类,
                oop_class: (B,)  预测类别
            }
        """
        with torch.no_grad():
            hmr_out = self.hmr(radar_input)
        
        # 拼接特征
        feat = torch.cat([
            hmr_out["pose"],
            hmr_out["shape"],
            hmr_out["trans"]
        ], dim=-1)
        
        oop_logits = self.classifier(feat)
        oop_class = oop_logits.argmax(dim=-1)
        
        return {
            **hmr_out,
            "oop_logits": oop_logits,
            "oop_class": oop_class
        }
    
    def post_to_oop_score(self, pose_3d: torch.Tensor) -> dict:
        """
        基于关节 3D 位置直接计算 OOP 风险分数
        
        Args:
            pose_3d: (B, 52, 3) 52 个关节的 3D 坐标
            
        Returns:
            scores: 各 OOP 风险分数
        """
        # 关键关节索引（SMPL-X 子集）
        HEAD = 12        # 头部
        SHOULDER_L = 2  # 左肩
        SHOULDER_R = 5  # 右肩
        HIP_L = 12      # 左髋
        HIP_R = 15      # 右髋
        KNEE_L = 18     # 左膝
        KNEE_R = 21     # 右膝
        ANKLE_L = 19    # 左踝
        ANKLE_R = 22    # 右踝
        
        # 计算姿态指标
        # 1. 前倾角度（躯干与垂直方向夹角）
        torso_vec = pose_3d[:, HIP_L] - pose_3d[:, SHOULDER_L]
        vertical = torch.tensor([0, 0, 1], device=pose_3d.device)
        forward_lean_angle = torch.acos(
            torch.clamp(
                torch.sum(torso_vec * vertical, dim=-1) / 
                (torch.norm(torso_vec, dim=-1) + 1e-8),
                -1, 1
            )
        )
        
        # 2. 侧倾角度（肩-肩连线与水平面夹角）
        shoulder_vec = pose_3d[:, SHOULDER_R] - pose_3d[:, SHOULDER_L]
        sideways_angle = torch.atan2(
            shoulder_vec[..., 2],
            torch.norm(shoulder_vec[..., :2], dim=-1)
        )
        
        # 3. 蜷缩检测（膝盖-髋距离）
        knee_hip_dist = torch.norm(
            pose_3d[:, KNEE_L] - pose_3d[:, HIP_L], dim=-1
        )
        fetal_score = torch.exp(-knee_hip_dist / 0.2)  # 距离越小分数越高
        
        # 4. 脚高度（踝关节 Z 坐标）
        feet_height = (pose_3d[:, ANKLE_L, 2] + pose_3d[:, ANKLE_R, 2]) / 2
        
        return {
            "forward_lean_deg": torch.degrees(forward_lean_angle),
            "sideways_lean_deg": torch.degrees(sideways_angle),
            "fetal_score": fetal_score,
            "feet_height_m": feet_height,
        }


# 测试代码
if __name__ == "__main__":
    model = OOPDetectionFromRadar()
    
    # 模拟雷达输入
    radar_input = torch.randn(2, 1, 32, 64, 128)  # (B, 1, angle, doppler, range)
    
    # 前向推理
    with torch.no_grad():
        output = model(radar_input)
    
    print("=== OOP 检测结果 ===")
    print(f"OOP 类别: {output['oop_class']}")
    for cls_idx, cls_name in model.OOP_CLASSES.items():
        prob = torch.softmax(output['oop_logits'], dim=-1)[0, cls_idx]
        print(f"  {cls_name}: {prob:.4f}")