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

class DualDomainFeatureConstructor(nn.Module):
    """
    双域特征构建
    
    输入：RAD张量 H_t (B, R, A, D)
        R=距离bins, A=角度bins, D=Doppler bins
    
    输出：
    - S_t: 空间幅度 (B, R, A) - Doppler维平均
    - V_t: Doppler特征 (B, R, A, D) - 原始Doppler签名
    """
    
    def __init__(self):
        super().__init__()
    
    def forward(self, H_t: torch.Tensor) -> tuple:
        """
        Args:
            H_t: (B, R, A, D) RAD张量
            
        Returns:
            S_t: (B, R, A) 空间幅度图
            V_t: (B, R, A, D) Doppler签名
        """
        # 空间幅度：沿Doppler维取平均（或最大值）
        S_t = H_t.mean(dim=-1)  # (B, R, A)
        
        # Doppler特征：保留原始Doppler维度
        V_t = H_t  # (B, R, A, D)
        
        return S_t, V_t


# 测试
if __name__ == "__main__":
    constructor = DualDomainFeatureConstructor()
    
    # 模拟RAD张量 (batch=1, range=64, angle=64, doppler=16)
    H_t = torch.randn(1, 64, 64, 16)
    
    S_t, V_t = constructor(H_t)
    
    print(f"输入 RAD: {H_t.shape}")
    print(f"空间幅度 S_t: {S_t.shape}")
    print(f"Doppler V_t: {V_t.shape}")

class RadarTokenizer(nn.Module):
    """
    空间和Doppler token化
    
    - 空间tokens: S_t分成 P_r × P_a 非重叠patch
    - Doppler tokens: V_t每个range-angle cell独立embed
    - 两组tokens共享相同网格
    """
    
    def __init__(self, range_bins: int = 64, angle_bins: int = 64,
                 doppler_bins: int = 16, patch_size: tuple = (4, 4),
                 embed_dim: int = 32):
        super().__init__()
        
        self.patch_size = patch_size
        self.embed_dim = embed_dim
        
        # 空间patch embedding
        self.spatial_embed = nn.Linear(
            patch_size[0] * patch_size[1], embed_dim
        )
        
        # Doppler embedding（每个cell独立embed）
        self.doppler_embed = nn.Linear(doppler_bins, embed_dim)
        
        # 位置编码
        n_patches_r = range_bins // patch_size[0]
        n_patches_a = angle_bins // patch_size[1]
        self.pos_embed = nn.Parameter(
            torch.randn(1, n_patches_r * n_patches_a, embed_dim)
        )
    
    def forward(self, S_t: torch.Tensor, V_t: torch.Tensor) -> tuple:
        """
        Args:
            S_t: (B, R, A) 空间幅度
            V_t: (B, R, A, D) Doppler签名
            
        Returns:
            spatial_tokens: (B, N, D) 空间tokens
            doppler_tokens: (B, R*A, D) Doppler tokens
        """
        B, R, A = S_t.shape
        
        # 空间token化：分patch
        pr, pa = self.patch_size
        S_patches = S_t.unfold(1, pr, pr).unfold(2, pa, pa)
        # (B, R/pr, A/pa, pr, pa)
        S_patches = S_patches.contiguous().view(B, -1, pr * pa)
        spatial_tokens = self.spatial_embed(S_patches)  # (B, N, D)
        
        # 添加位置编码
        spatial_tokens = spatial_tokens + self.pos_embed
        
        # Doppler token化：每个cell独立embed
        V_flat = V_t.view(B, R * A, -1)  # (B, R*A, D)
        doppler_tokens = self.doppler_embed(V_flat)  # (B, R*A, D)
        
        return spatial_tokens, doppler_tokens


# 测试
if __name__ == "__main__":
    tokenizer = RadarTokenizer(
        range_bins=64, angle_bins=64, doppler_bins=16,
        patch_size=(4, 4), embed_dim=32
    )
    
    S_t = torch.randn(1, 64, 64)
    V_t = torch.randn(1, 64, 64, 16)
    
    spatial_tokens, doppler_tokens = tokenizer(S_t, V_t)
    
    print(f"空间tokens: {spatial_tokens.shape} (B, N_patches, D)")
    print(f"Doppler tokens: {doppler_tokens.shape} (B, R*A, D)")
    print(f"N_patches = (64/4) × (64/4) = {16*16}")

class ConfidenceGatedCrossAttention(nn.Module):
    """
    置信度门控局部交叉注意力
    
    核心设计：
    1. 可学习门 g_{t,j} ∈ [0,1] 评估每个Doppler token的运动相关性
    2. 空间token只关注局部邻域 N(i) 内的Doppler tokens
    3. 门控权重调制注意力，抑制不可靠Doppler响应
    
    这解决了两个关键问题：
    - 跨区域频谱泄漏：局部注意力限制
    - 杂波/伪影Doppler：置信度门控
    """
    
    def __init__(self, embed_dim: int = 32, neighborhood: int = 3):
        super().__init__()
        self.embed_dim = embed_dim
        self.neighborhood = neighborhood
        
        # 置信度门（评估Doppler token质量）
        self.gate = nn.Sequential(
            nn.Linear(embed_dim, embed_dim // 2),
            nn.GELU(),
            nn.Linear(embed_dim // 2, 1),
            nn.Sigmoid()
        )
        
        # 交叉注意力
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        
        self.scale = embed_dim ** -0.5
    
    def forward(self, spatial_tokens: torch.Tensor,
                doppler_tokens: torch.Tensor,
                n_patches_r: int, n_patches_a: int) -> torch.Tensor:
        """
        Args:
            spatial_tokens: (B, N, D) N = n_patches_r * n_patches_a
            doppler_tokens: (B, R*A, D) 每个cell一个Doppler token
            n_patches_r, n_patches_a: patch网格尺寸
            
        Returns:
            conditioned_tokens: (B, N, D) 运动条件化的空间tokens
        """
        B, N, D = spatial_tokens.shape
        R_A = doppler_tokens.shape[1]
        
        # 1. 计算Doppler置信度门
        gates = self.gate(doppler_tokens)  # (B, R*A, 1)
        
        # 2. 将Doppler tokens重塑为2D网格
        # 假设 R*A 可以映射到 patch 网格
        # 每个patch对应多个Doppler cells
        cells_per_patch_r = R_A // (n_patches_r * n_patches_a) // n_patches_a + 1
        
        # 简化：将Doppler tokens下采样到patch网格
        # 实际实现中使用更精确的映射
        doppler_2d = doppler_tokens.view(B, n_patches_r, -1, D)
        gates_2d = gates.view(B, n_patches_r, -1, 1)
        
        # 3. 局部交叉注意力
        Q = self.q_proj(spatial_tokens)  # (B, N, D)
        
        # 为每个空间token，计算对局部Doppler tokens的注意力
        # 简化实现：使用全局注意力（实际使用局部窗口）
        K = self.k_proj(doppler_tokens)  # (B, R*A, D)
        V = self.v_proj(doppler_tokens)  # (B, R*A, D)
        
        # 注意力分数
        attn = torch.bmm(Q, K.transpose(1, 2)) * self.scale  # (B, N, R*A)
        
        # 门控调制
        attn = attn * gates.transpose(1, 2)  # 广播门控
        
        # 局部掩码（只允许3x3邻域）
        # 实际实现使用稀疏注意力或窗口注意力
        # 这里简化为全局注意力 + 门控
        
        attn = torch.softmax(attn, dim=-1)
        
        # 加权求和
        out = torch.bmm(attn, V)  # (B, N, D)
        
        # 残差连接
        conditioned = spatial_tokens + out
        
        return conditioned


class PULSEPrompting(nn.Module):
    """
    PULSE 完整模块
    
    可作为插件替换现有backbone的前端融合
    """
    
    def __init__(self, range_bins=64, angle_bins=64, doppler_bins=16,
                 patch_size=(4, 4), embed_dim=32,
                 neighborhood=3, beta=1.0):
        super().__init__()
        
        self.constructor = DualDomainFeatureConstructor()
        self.tokenizer = RadarTokenizer(
            range_bins, angle_bins, doppler_bins,
            patch_size, embed_dim
        )
        self.cross_attn = ConfidenceGatedCrossAttention(
            embed_dim, neighborhood
        )
        
        self.n_patches_r = range_bins // patch_size[0]
        self.n_patches_a = angle_bins // patch_size[1]
        
        # 输出投影
        self.output_proj = nn.Linear(embed_dim, embed_dim)
    
    def forward(self, H_t: torch.Tensor) -> torch.Tensor:
        """
        Args:
            H_t: (B, R, A, D) RAD张量
            或 (B, T, R, A, D) 多帧序列
        
        Returns:
            spatial_tokens: (B, N, D) 运动条件化空间tokens
        """
        if H_t.dim() == 5:
            # 多帧模式
            B, T, R, A, D = H_t.shape
            # 简化：取最后一帧（实际用置信度加权聚合）
            H_t = H_t[:, -1]
        
        S_t, V_t = self.constructor(H_t)
        spatial_tokens, doppler_tokens = self.tokenizer(S_t, V_t)
        conditioned = self.cross_attn(
            spatial_tokens, doppler_tokens,
            self.n_patches_r, self.n_patches_a
        )
        
        return self.output_proj(conditioned)


# 测试
if __name__ == "__main__":
    pulse = PULSEPrompting(
        range_bins=64, angle_bins=64, doppler_bins=16,
        patch_size=(4, 4), embed_dim=32,
        neighborhood=3, beta=1.0
    )
    
    # 模拟RAD张量
    H_t = torch.randn(1, 64, 64, 16)
    
    output = pulse(H_t)
    
    print("=== PULSE 模块 ===")
    print(f"输入 RAD: {H_t.shape}")
    print(f"输出 tokens: {output.shape}")
    
    total_params = sum(p.numel() for p in pulse.parameters())
    print(f"参数量: {total_params:,}")
    print(f"模型大小: {total_params * 4 / 1024:.1f} KB")

class PULSEMultiFrame(nn.Module):
    """
    PULSE-KF: 多帧扩展
    
    用置信度加权的K帧Doppler聚合替代单帧Doppler tokens
    不修改空间backbone
    """
    
    def __init__(self, k_frames: int = 9, **kwargs):
        super().__init__()
        self.k_frames = k_frames
        self.pulse = PULSEPrompting(**kwargs)
        
        # 时序置信度聚合
        self.temporal_gate = nn.Sequential(
            nn.Linear(kwargs.get('embed_dim', 32), 1),
            nn.Sigmoid()
        )
    
    def forward(self, H_window: torch.Tensor) -> torch.Tensor:
        """
        Args:
            H_window: (B, K, R, A, D) K帧RAD张量序列
        Returns:
            spatial_tokens: (B, N, D)
        """
        B, K, R, A, D = H_window.shape
        
        # 处理每帧
        all_doppler = []
        for k in range(K):
            H_t = H_window[:, k]
            S_t, V_t = self.pulse.constructor(H_t)
            _, doppler_tokens = self.pulse.tokenizer(S_t, V_t)
            all_doppler.append(doppler_tokens)
        
        # 时序置信度加权聚合
        stacked = torch.stack(all_doppler, dim=1)  # (B, K, R*A, D)
        
        weights = self.temporal_gate(stacked)  # (B, K, R*A, 1)
        weighted = (stacked * weights).sum(dim=1) / (weights.sum(dim=1) + 1e-8)
        
        # 用聚合后的Doppler tokens
        H_last = H_window[:, -1]
        S_t, _ = self.pulse.constructor(H_last)
        spatial_tokens, _ = self.pulse.tokenizer(S_t, weighted)
        
        conditioned = self.pulse.cross_attn(
            spatial_tokens, weighted,
            self.pulse.n_patches_r, self.pulse.n_patches_a
        )
        
        return self.pulse.output_proj(conditioned)


# 测试
if __name__ == "__main__":
    pulse_kf = PULSEMultiFrame(k_frames=9, range_bins=64, angle_bins=64,
                                doppler_bins=16, patch_size=(4, 4),
                                embed_dim=32, neighborhood=3)
    
    # 模拟9帧序列
    H_window = torch.randn(1, 9, 64, 64, 16)
    output = pulse_kf(H_window)
    
    print(f"输入 9帧RAD: {H_window.shape}")
    print(f"输出 tokens: {output.shape}")
    
    total_params = sum(p.numel() for p in pulse_kf.parameters())
    print(f"参数量: {total_params:,}")

class CabinOOPDetector:
    """
    基于 PULSE + mmWave 雷达的座舱OOP检测器
    
    优势：
    1. 隐私保护（无图像）
    2. 穿透毯子/衣物
    3. 5.1ms推理延迟（实时）
    4. 12M参数（可部署边缘NPU）
    5. PULSE降低了MPJVE（抖动），适合时序稳定的OOP判断
    
    部署方案：
    - 雷达：Infineon XENSIV 60GHz 或 TI IWR6843
    - 处理器：QCS8255 Hexagon NPU 或 STM32N6
    """
    
    def __init__(self, threshold_lean: float = 30.0,
                 threshold_kneel: float = 0.7):
        self.threshold_lean = threshold_lean  # 前倾角度阈值
        self.threshold_kneel = threshold_kneel  # 跪立高度阈值
    
    def detect_oop_from_joints(self, joints_3d: np.ndarray) -> dict:
        """
        从3D关节坐标检测OOP
        
        Args:
            joints_3d: (J, 3) 3D关节坐标 (mm)
                       J=17 or 24 depending on skeleton format
        Returns:
            OOP检测结果
        """
        # 关键关节（简化索引）
        PELVIS, SPINE, NECK, HEAD = 0, 3, 12, 15
        L_SHOULDER, R_SHOULDER = 16, 17
        L_HIP, R_HIP = 1, 2
        L_KNEE, R_KNEE = 4, 5
        L_ANKLE, R_ANKLE = 7, 8
        
        # 1. 前倾检测：躯干向量与垂直方向夹角
        torso_vec = joints_3d[NECK] - joints_3d[PELVIS]
        forward_angle = np.degrees(np.arctan2(
            abs(torso_vec[2]),  # z: 前后
            max(abs(torso_vec[1]), 1.0)  # y: 上下
        ))
        
        # 2. 侧倾检测
        side_angle = np.degrees(np.arctan2(
            torso_vec[0],  # x: 左右
            max(abs(torso_vec[1]), 1.0)
        ))
        
        # 3. 跪立检测：膝盖高度
        knee_y = min(joints_3d[L_KNEE][1], joints_3d[R_KNEE][1])
        hip_y = max(joints_3d[L_HIP][1], joints_3d[R_HIP][1])
        is_kneeling = knee_y > hip_y * 0.5  # 膝盖接近臀部高度
        
        results = {}
        if forward_angle > self.threshold_lean:
            results['forward_lean'] = {
                'angle': float(forward_angle),
                'severity': min(forward_angle / 60, 1.0)
            }
        
        if abs(side_angle) > self.threshold_lean:
            results['side_lean'] = {
                'angle': float(side_angle),
                'severity': min(abs(side_angle) / 45, 1.0)
            }
        
        if is_kneeling:
            results['kneeling'] = {
                'knee_height_ratio': float(knee_y / (hip_y + 1e-6)),
                'severity': 0.8
            }
        
        return {
            'is_oop': len(results) > 0,
            'oop_types': list(results.keys()),
            'details': results,
            'joint_metrics': {
                'forward_angle': float(forward_angle),
                'side_angle': float(side_angle),
                'knee_hip_ratio': float(knee_y / (hip_y + 1e-6))
            }
        }


# 测试
if __name__ == "__main__":
    detector = CabinOOPDetector()
    
    # 模拟正常坐姿关节
    normal_joints = np.array([
        [0, -400, 0],     # pelvis
        [-100, -350, 0],  # left_hip
        [100, -350, 0],   # right_hip
        [0, -300, 0],     # spine1
        [-100, -200, 0],  # left_knee
        [100, -200, 0],   # right_knee
        [0, -250, 0],     # spine2
        [-100, -50, 0],   # left_ankle
        [100, -50, 0],    # right_ankle
        [0, -200, 0],     # spine3
        [-100, -50, 0],   # left_foot
        [100, -50, 0],    # right_foot
        [0, -150, 0],     # neck
        [-150, -180, 0],  # left_collar
        [150, -180, 0],   # right_collar
        [0, -100, 0],     # head
        [-200, -200, 0],  # left_shoulder
        [200, -200, 0],   # right_shoulder
        [-300, -300, 0],  # left_elbow
        [300, -300, 0],   # right_elbow
        [-400, -380, 0],  # left_wrist
        [400, -380, 0],   # right_wrist
        [-450, -380, 0],  # left_hand
        [450, -380, 0],   # right_hand
    ], dtype=np.float32)
    
    # 模拟前倾坐姿
    lean_joints = normal_joints.copy()
    lean_joints[NECK := 12] = [0, -100, 50]  # neck前倾
    lean_joints[HEAD := 15] = [0, -50, 60]   # head前倾
    
    # 模拟跪立
    kneel_joints = normal_joints.copy()
    kneel_joints[4] = [-100, -350, 0]  # left_knee抬高
    kneel_joints[5] = [100, -350, 0]   # right_knee抬高
    
    print("=== 正常坐姿 ===")
    print(detector.detect_oop_from_joints(normal_joints))
    
    print("\n=== 前倾OOP ===")
    print(detector.detect_oop_from_joints(lean_joints))
    
    print("\n=== 跪立OOP ===")
    print(detector.detect_oop_from_joints(kneel_joints))

# 集成示例（来自论文代码）
from pulse import PULSEPrompting

# 替换现有backbone的前端
pulse_front_end = PULSEPrompting(
    range_bins=64, angle_bins=64, doppler_bins=16,
    patch_size=(4, 4), embed_dim=32,
    neighborhood=3, beta=1.0
)

# 替换前端融合
spatial_tokens = pulse_front_end(H_t)  # 输入RAD张量
# 然后送入现有backbone
output = existing_backbone(spatial_tokens)