# SMPL-X 参数结构
smpl_x_params = {
    "shape": (10,),        # 形状参数 β，控制体型
    "pose": (55, 3),       # 姿态参数 θ，55个关节×3轴旋转
    "expression": (10,),   # 表情参数（网格恢复中通常忽略）
    "trans": (3,),         # 全局平移
    "gender": "neutral"   # 性别（影响形状空间）
}

# 总参数维度：10 + 165 + 3 = 178（核心参数）
# 输出网格：10475 个顶点，52 个关节

import torch
import torch.nn as nn

class FullVolumeLocator(nn.Module):
    """
    从完整雷达体积中定位人体中心
    
    输入：雷达张量 (B, C, D_range, D_doppler, D_angle)
    输出：人体中心坐标 (B, 3) + 全局特征 (B, C', D_r, D_d, D_a)
    """
    def __init__(self, in_channels: int, hidden: int = 64):
        super().__init__()
        # 3D 卷积保持空间结构
        self.conv3d = nn.Sequential(
            nn.Conv3d(in_channels, hidden, kernel_size=3, padding=1),
            nn.BatchNorm3d(hidden),
            nn.ReLU(),
            nn.Conv3d(hidden, hidden * 2, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm3d(hidden * 2),
            nn.ReLU(),
            nn.AdaptiveAvgPool3d(1)
        )
        # 中心回归头
        self.center_head = nn.Linear(hidden * 2, 3)
        
    def forward(self, x):
        # x: (B, C, D_r, D_d, D_a)
        features = self.conv3d(x)  # (B, 2*hidden, 1, 1, 1)
        features = features.flatten(1)
        center = self.center_head(features)  # (B, 3)
        return center, features

class CenterGuidedROI(nn.Module):
    """
    根据定位中心提取人体局部雷达体积
    
    避免 FPN 的全局处理稀释弱肢体反射证据
    """
    def __init__(self, roi_size=(32, 16, 32)):
        super().__init__()
        self.roi_size = roi_size  # (range, doppler, angle) 裁剪大小
        
    def forward(self, x: torch.Tensor, center: torch.Tensor):
        """
        Args:
            x: 完整雷达张量 (B, C, D_r, D_d, D_a)
            center: 人体中心 (B, 3) -> (range_idx, doppler_idx, angle_idx)
        Returns:
            local_volume: 人体中心局部体积 (B, C, roi_r, roi_d, roi_a)
        """
        B = x.shape[0]
        local_volumes = []
        for i in range(B):
            cr, cd, ca = center[i].long()
            r_start = max(0, cr - self.roi_size[0] // 2)
            d_start = max(0, cd - self.roi_size[1] // 2)
            a_start = max(0, ca - self.roi_size[2] // 2)
            
            r_end = min(x.shape[2], r_start + self.roi_size[0])
            d_end = min(x.shape[3], d_start + self.roi_size[1])
            a_end = min(x.shape[4], a_start + self.roi_size[2])
            
            local_volumes.append(x[i, :, r_start:r_end, d_start:d_end, a_start:a_end])
        
        return torch.stack(local_volumes, dim=0)

class GlobalLocalFusion(nn.Module):
    """
    门控融合全局定位特征与局部人体特征
    """
    def __init__(self, global_dim: int, local_dim: int, out_dim: int):
        super().__init__()
        self.global_proj = nn.Linear(global_dim, out_dim)
        self.local_proj = nn.Conv3d(local_dim, out_dim, kernel_size=1)
        self.gate = nn.Sequential(
            nn.Linear(out_dim * 2, out_dim),
            nn.Sigmoid()
        )
        self.output_proj = nn.Conv3d(out_dim, out_dim, kernel_size=3, padding=1)
        
    def forward(self, global_feat, local_volume):
        # global_feat: (B, global_dim)
        # local_volume: (B, local_dim, roi_r, roi_d, roi_a)
        g = self.global_proj(global_feat)  # (B, out_dim)
        l = self.local_proj(local_volume)  # (B, out_dim, roi_r, roi_d, roi_a)
        
        # 广播全局特征到局部体积
        g_expanded = g.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
        g_expanded = g_expanded.expand(-1, -1, l.shape[2], l.shape[3], l.shape[4])
        
        # 门控融合
        gate_weights = self.gate(
            torch.cat([g_expanded, l], dim=1).mean(dim=(2, 3, 4))
        ).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
        
        fused = gate_weights * g_expanded + (1 - gate_weights) * l
        return self.output_proj(fused)  # RaSS token

class DualStateTemporalModule(nn.Module):
    """
    双状态因果时间模块
    
    慢速状态：积累形态证据（shape, gender），使用 EMA 更新
    快速状态：跟踪运动状态（pose, translation），使用因果循环
    """
    def __init__(self, token_dim: int, morphology_dim: int = 10, motion_dim: int = 168):
        super().__init__()
        self.morphology_dim = morphology_dim  # shape(10) + gender(1) ≈ 11
        self.motion_dim = motion_dim          # pose(165) + trans(3) = 168
        
        # 慢速形态状态
        self.morphology_gate = nn.Linear(token_dim, morphology_dim)
        self.morphology_ema = None  # 推理时的持久状态
        
        # 快速运动状态
        self.motion_rnn = nn.GRUCell(token_dim, motion_dim)
        self.motion_head = nn.Linear(motion_dim, motion_dim)
        
    def forward(self, z_t: torch.Tensor, prev_motion_state: torch.Tensor = None):
        """
        Args:
            z_t: RaSS token at time t, shape (B, token_dim)
            prev_motion_state: 上一步运动状态 (B, motion_dim)
        Returns:
            morphology_out: 形态参数 (B, morphology_dim)
            motion_out: 运动参数 (B, motion_dim)
            new_motion_state: 更新后的运动状态
        """
        B = z_t.shape[0]
        
        # 慢速状态：token-conditioned EMA
        gate_val = torch.sigmoid(self.morphology_gate(z_t))  # (B, morph_dim)
        if self.morphology_ema is None:
            self.morphology_ema = torch.zeros(B, self.morphology_dim, device=z_t.device)
        
        # 弱帧降权：gate值低时保持历史状态
        self.morphology_ema = gate_val * z_t[:, :self.morphology_dim] + \
                              (1 - gate_val) * self.morphology_ema
        
        # 快速状态：因果 RNN
        if prev_motion_state is None:
            prev_motion_state = torch.zeros(B, self.motion_dim, device=z_t.device)
        new_motion_state = self.motion_rnn(z_t, prev_motion_state)
        motion_out = self.motion_head(new_motion_state)
        
        return self.morphology_ema, motion_out, new_motion_state

class StreamingConfig:
    """RaStream 流式推理配置"""
    def __init__(self, Tw=3, T=8, s=1, warmup=2):
        self.Tw = Tw          # 每步雷达观测帧数
        self.T = T             # 时间展开步数
        self.s = s             # 滑动步长
        self.warmup = warmup   # 预热步数（重放前几帧）
        
    def compute_output_rate(self, radar_fps=30):
        """计算输出帧率"""
        return radar_fps / self.s
    
    def compute_memory(self, token_dim=256):
        """估算流式内存（不含模型权重）"""
        # 慢速状态 + 快速状态 + 当前 token
        return (token_dim +  # morphology EMA
                token_dim +  # motion RNN hidden
                token_dim)   # current token
    
    def compute_latency_budget(self, max_latency_ms=33):
        """计算单帧延迟预算（30fps = 33ms）"""
        per_step_budget = max_latency_ms * self.s
        return per_step_budget

# 典型配置
config = StreamingConfig(Tw=3, T=8, s=1, warmup=2)
print(f"输出帧率: {config.compute_output_rate(30):.1f} fps")
print(f"流式内存: {config.compute_memory(256)} floats = {config.compute_memory(256) * 4 / 1024:.1f} KB")

"""
RaStream 推理管道复现
论文：RaStream: Edge-Deployable Streaming Human Mesh Recovery from mmWave Radar
arXiv: 2608.11791
"""

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

@dataclass
class RadarConfig:
    """mmWave 雷达配置"""
    range_bins: int = 64       # 距离维
    doppler_bins: int = 16     # 多普勒维
    angle_bins: int = 32       # 角度维
    channels: int = 4          # 天线通道

@dataclass 
class SMPLXConfig:
    """SMPL-X 模型配置"""
    shape_dim: int = 10        # β 参数
    pose_dim: int = 165        # 55 * 3 关节角度
    trans_dim: int = 3         # 全局平移
    gender_dim: int = 1        # 性别
    total_dim: int = 179       # 总参数维度


class RaSSSpatialEncoder(nn.Module):
    """
    RaSS: Radar-aware Spatial Structure Encoder
    论文 Section II-A
    
    功能：
    1. 保留 3D 雷达体积结构
    2. 全体积人体定位
    3. 中心引导 RoI 提取
    4. 全局-局部融合
    """
    def __init__(self, radar_config: RadarConfig, hidden: int = 64, token_dim: int = 256):
        super().__init__()
        self.radar_config = radar_config
        self.token_dim = token_dim
        
        # 全体积定位分支
        self.locator_conv = nn.Sequential(
            nn.Conv3d(radar_config.channels, hidden, 3, padding=1),
            nn.BatchNorm3d(hidden),
            nn.ReLU(inplace=True),
            nn.Conv3d(hidden, hidden * 2, 3, stride=2, padding=1),
            nn.BatchNorm3d(hidden * 2),
            nn.ReLU(inplace=True),
            nn.Conv3d(hidden * 2, hidden * 4, 3, stride=2, padding=1),
            nn.BatchNorm3d(hidden * 4),
            nn.ReLU(inplace=True),
        )
        self.locator_head = nn.Sequential(
            nn.AdaptiveAvgPool3d(1),
            nn.Flatten(),
            nn.Linear(hidden * 4, hidden * 2),
            nn.ReLU(),
            nn.Linear(hidden * 2, 3),
            nn.Sigmoid()  # 归一化到 [0, 1] 范围
        )
        
        # 局部体积编码
        self.local_encoder = nn.Sequential(
            nn.Conv3d(radar_config.channels, hidden, 3, padding=1),
            nn.BatchNorm3d(hidden),
            nn.ReLU(inplace=True),
            nn.Conv3d(hidden, hidden * 2, 3, stride=2, padding=1),
            nn.BatchNorm3d(hidden * 2),
            nn.ReLU(inplace=True),
            nn.Conv3d(hidden * 2, hidden * 4, 3, stride=2, padding=1),
            nn.BatchNorm3d(hidden * 4),
            nn.ReLU(inplace=True),
        )
        
        # 全局-局部门控融合
        self.global_proj = nn.Linear(hidden * 4, token_dim)
        self.local_proj = nn.Linear(hidden * 4, token_dim)
        self.gate = nn.Sequential(
            nn.Linear(token_dim * 2, token_dim),
            nn.Sigmoid()
        )
        
    def forward(self, radar_tensor: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            radar_tensor: (B, C, D_r, D_d, D_a)
        Returns:
            token: (B, token_dim) - RaSS 空间 token
            center: (B, 3) - 归一化人体中心
        """
        B = radar_tensor.shape[0]
        
        # 全体积定位
        global_feat = self.locator_conv(radar_tensor)  # (B, 4*hidden, d_r, d_d, d_a)
        center = self.locator_head(global_feat)  # (B, 3)
        
        # RoI 提取（简化版：用中心裁剪固定大小区域）
        cr = (center[:, 0] * radar_tensor.shape[2]).long().clamp(8, radar_tensor.shape[2] - 8)
        cd = (center[:, 1] * radar_tensor.shape[3]).long().clamp(4, radar_tensor.shape[3] - 4)
        ca = (center[:, 2] * radar_tensor.shape[4]).long().clamp(8, radar_tensor.shape[4] - 8)
        
        local_volumes = []
        for i in range(B):
            r_s, d_s, a_s = int(cr[i]) - 8, int(cd[i]) - 4, int(ca[i]) - 8
            r_e, d_e, a_e = r_s + 16, d_s + 8, a_s + 16
            local_volumes.append(radar_tensor[i, :, r_s:r_e, d_s:d_e, a_s:a_e])
        local_volume = torch.stack(local_volumes, dim=0)
        
        # 局部编码
        local_feat = self.local_encoder(local_volume)
        local_feat = local_feat.flatten(2).mean(dim=2)  # 全局平均池化
        
        # 全局-局部融合
        global_feat_flat = global_feat.flatten(2).mean(dim=2)
        g = self.global_proj(global_feat_flat)
        l = self.local_proj(local_feat)
        
        gate_weight = self.gate(torch.cat([g, l], dim=1))
        token = gate_weight * g + (1 - gate_weight) * l
        
        return token, center


class DualStateTemporalRefinement(nn.Module):
    """
    双状态因果时间精炼模块
    论文 Section III
    
    慢速状态：身体形态（shape + gender）通过 EMA 累积
    快速状态：运动（pose + translation）通过 GRU 跟踪
    """
    def __init__(self, token_dim: int = 256, smplx_config: SMPLXConfig = SMPLXConfig()):
        super().__init__()
        self.smplx_config = smplx_config
        
        # 慢速形态分支
        morph_dim = smplx_config.shape_dim + smplx_config.gender_dim  # 11
        self.morph_gate = nn.Linear(token_dim, morph_dim)
        self.morph_proj = nn.Linear(token_dim, morph_dim)
        self.morph_out = nn.Linear(morph_dim, morph_dim)
        
        # 快速运动分支
        motion_dim = smplx_config.pose_dim + smplx_config.trans_dim  # 168
        self.motion_gru = nn.GRUCell(token_dim, motion_dim)
        self.motion_out = nn.Sequential(
            nn.Linear(motion_dim, motion_dim),
            nn.ReLU(),
            nn.Linear(motion_dim, motion_dim)
        )
        
        # 6D 旋转表示（连续旋转，避免万向锁）
        self.rot_6d_dim = smplx_config.pose_dim // 3 * 6  # 55 * 6 = 330
        
    def forward(self, z_t: torch.Tensor, 
               prev_morph_state: Optional[torch.Tensor] = None,
               prev_motion_state: Optional[torch.Tensor] = None) -> Tuple:
        """
        Args:
            z_t: 当前时刻 RaSS token (B, token_dim)
            prev_morph_state: 上一步形态状态
            prev_motion_state: 上一步运动状态
        Returns:
            morph_params, motion_params, new_morph_state, new_motion_state
        """
        B = z_t.shape[0]
        morph_dim = self.smplx_config.shape_dim + self.smplx_config.gender_dim
        motion_dim = self.smplx_config.pose_dim + self.smplx_config.trans_dim
        
        # 慢速形态状态：token-conditioned EMA
        gate_val = torch.sigmoid(self.morph_gate(z_t))  # (B, morph_dim)
        morph_candidate = self.morph_proj(z_t)  # (B, morph_dim)
        
        if prev_morph_state is None:
            prev_morph_state = torch.zeros(B, morph_dim, device=z_t.device)
        
        # 弱帧降权：gate值低时保持历史状态
        new_morph_state = gate_val * morph_candidate + (1 - gate_val) * prev_morph_state
        morph_params = self.morph_out(new_morph_state)
        
        # 快速运动状态：因果 GRU
        if prev_motion_state is None:
            prev_motion_state = torch.zeros(B, motion_dim, device=z_t.device)
        
        new_motion_state = self.motion_gru(z_t, prev_motion_state)
        motion_params = self.motion_out(new_motion_state)
        
        # 分离 pose 和 trans
        pose_params = motion_params[:, :self.smplx_config.pose_dim]
        trans_params = motion_params[:, self.smplx_config.pose_dim:]
        
        # 形状和性别
        shape_params = morph_params[:, :self.smplx_config.shape_dim]
        gender_params = morph_params[:, self.smplx_config.shape_dim:]
        
        return {
            'shape': shape_params,      # (B, 10)
            'gender': gender_params,     # (B, 1)
            'pose': pose_params,         # (B, 165)
            'trans': trans_params,       # (B, 3)
        }, new_morph_state, new_motion_state


class RaStream(nn.Module):
    """
    RaStream: 完整的流式 mmWave 人体网格恢复模型
    
    论文：arXiv:2608.11791
    应用场景：IMS OOP 异常姿态检测、座舱乘员 3D 感知
    """
    def __init__(self, radar_config: RadarConfig = RadarConfig(), 
                 smplx_config: SMPLXConfig = SMPLXConfig(),
                 token_dim: int = 256):
        super().__init__()
        self.spatial_encoder = RaSSSpatialEncoder(radar_config, hidden=64, token_dim=token_dim)
        self.temporal_module = DualStateTemporalRefinement(token_dim, smplx_config)
        
    def forward(self, radar_stream: torch.Tensor, 
               states: Optional[Tuple] = None) -> Tuple:
        """
        流式推理
        
        Args:
            radar_stream: (B, T, C, D_r, D_d, D_a) 雷达流
            states: (morph_state, motion_state) 前一时刻状态
        Returns:
            all_params: list of param dicts
            final_states: (morph_state, motion_state)
        """
        B, T = radar_stream.shape[:2]
        all_params = []
        
        morph_state = None
        motion_state = None
        if states is not None:
            morph_state, motion_state = states
        
        for t in range(T):
            # 空间编码
            z_t, center = self.spatial_encoder(radar_stream[:, t])
            
            # 时间精炼
            params, morph_state, motion_state = self.temporal_module(
                z_t, morph_state, motion_state
            )
            
            params['center'] = center
            all_params.append(params)
        
        return all_params, (morph_state, motion_state)


# ==================== 测试与验证 ====================

def test_rastream():
    """RaStream 模型测试"""
    radar_config = RadarConfig(range_bins=64, doppler_bins=16, angle_bins=32, channels=4)
    smplx_config = SMPLXConfig()
    
    model = RaStream(radar_config, smplx_config, token_dim=256)
    
    # 模拟雷达流：(batch=2, time=10, channels=4, range=64, doppler=16, angle=32)
    radar_stream = torch.randn(2, 10, 4, 64, 16, 32) * 10  # 模拟雷达回波
    
    # 流式推理
    all_params, final_states = model(radar_stream)
    
    # 输出分析
    print("=" * 60)
    print("RaStream 流式推理测试结果")
    print("=" * 60)
    print(f"输入: batch=2, time=10, channels=4, range=64, doppler=16, angle=32")
    print(f"输出参数数: {len(all_params)} (时间步)")
    print()
    
    # 检查最后一步参数
    final_params = all_params[-1]
    for key, val in final_params.items():
        print(f"  {key}: shape={val.shape}, range=[{val.min():.3f}, {val.max():.3f}]")
    
    # 参数量统计
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n模型参数量: {total_params / 1e6:.2f}M")
    
    # 内存估算
    radar_input_size = 2 * 10 * 4 * 64 * 16 * 32 * 4 / 1024 / 1024
    print(f"雷达输入大小: {radar_input_size:.1f} MB")
    
    return model, all_params


def test_streaming_stability():
    """测试流式推理的时间稳定性"""
    print("\n" + "=" * 60)
    print("时间稳定性测试")
    print("=" * 60)
    
    model = RaStream()
    
    # 模拟 50 帧雷达流
    T = 50
    radar_stream = torch.randn(1, T, 4, 64, 16, 32) * 10
    
    # 加入一个弱帧（低信噪比）
    radar_stream[0, 25] *= 0.1  # 第25帧为弱帧
    
    all_params, _ = model(radar_stream)
    
    # 计算帧间变化
    shape_changes = []
    pose_changes = []
    for t in range(1, T):
        shape_diff = (all_params[t]['shape'] - all_params[t-1]['shape']).norm().item()
        pose_diff = (all_params[t]['pose'] - all_params[t-1]['pose']).norm().item()
        shape_changes.append(shape_diff)
        pose_changes.append(pose_diff)
    
    print(f"形状参数帧间变化: 均值={np.mean(shape_changes):.4f}, 峰值={np.max(shape_changes):.4f}")
    print(f"姿态参数帧间变化: 均值={np.mean(pose_changes):.4f}, 峰值={np.max(pose_changes):.4f}")
    print(f"弱帧(第25帧)形状变化: {shape_changes[24]:.4f} (应低于平均)")
    print(f"弱帧(第25帧)姿态变化: {pose_changes[24]:.4f} (可能高于平均，因快速状态响应)")


def test_edge_deployment_profile():
    """边缘部署性能分析"""
    print("\n" + "=" * 60)
    print("边缘部署性能分析（模拟 Jetson Orin Nano）")
    print("=" * 60)
    
    model = RaStream()
    model.eval()
    
    # 单帧推理
    single_frame = torch.randn(1, 1, 4, 64, 16, 32)
    
    # 计时
    import time
    with torch.no_grad():
        # 预热
        for _ in range(5):
            _ = model(single_frame)
        
        # 正式计时
        times = []
        for _ in range(20):
            start = time.perf_counter()
            _ = model(single_frame)
            times.append((time.perf_counter() - start) * 1000)
    
    print(f"单帧推理延迟: {np.mean(times):.2f} ms (±{np.std(times):.2f})")
    print(f"理论帧率: {1000 / np.mean(times):.1f} fps")
    
    # 内存
    param_size = sum(p.numel() * p.element_size() for p in model.parameters())
    print(f"模型大小: {param_size / 1024 / 1024:.2f} MB")
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"参数量: {total_params / 1e6:.2f}M")
    
    # 流式状态内存
    token_dim = 256
    morph_dim = 11
    motion_dim = 168
    state_memory = (morph_dim + motion_dim + token_dim) * 4  # float32
    print(f"流式状态内存: {state_memory} bytes = {state_memory / 1024:.2f} KB")


if __name__ == "__main__":
    model, params = test_rastream()
    test_streaming_stability()
    test_edge_deployment_profile()