import torch
import torch.nn as nn
import torch.nn.functional as F

class EdgeConv(nn.Module):
    """
    EdgeConv层 - DGCNN的核心操作
    
    对每个点的k近邻计算边特征，再通过MLP聚合
    """
    def __init__(self, in_channels: int, out_channels: int, k: int = 16):
        super().__init__()
        self.k = k
        # 边特征变换MLP
        self.mlp = nn.Sequential(
            nn.Linear(in_channels * 2, out_channels),
            nn.BatchNorm1d(out_channels),
            nn.LeakyReLU(0.2),
            nn.Linear(out_channels, out_channels),
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 点云数据, shape=(B, N, C) 
                B=batch_size, N=点数, C=特征维度
        Returns:
            out: 变换后的特征, shape=(B, N, out_channels)
        """
        B, N, C = x.shape
        
        # 计算成对距离矩阵
        dist_matrix = torch.cdist(x, x)  # (B, N, N)
        
        # 找到k近邻
        _, knn_idx = torch.topk(-dist_matrix, self.k, dim=-1)  # (B, N, k)
        
        # 收集邻居特征
        knn_features = torch.gather(
            x.unsqueeze(2).expand(-1, -1, self.k, -1),
            1,
            knn_idx.unsqueeze(-1).expand(-1, -1, -1, C)
        )  # (B, N, k, C)
        
        # 计算边特征: 中心点特征 + 差值特征
        center_feature = x.unsqueeze(2).expand(-1, -1, self.k, -1)
        edge_feature = torch.cat([
            center_feature,  # 中心点
            knn_features - center_feature  # 相对偏移
        ], dim=-1)  # (B, N, k, 2C)
        
        # MLP变换 + max聚合
        edge_features = self.mlp(edge_features.reshape(B * N * self.k, -1))
        edge_features = edge_features.reshape(B, N, self.k, -1)
        out = edge_features.max(dim=2)[0]  # (B, N, out_channels)
        
        return out


class DGCNNFeatureExtractor(nn.Module):
    """
    DGCNN空间特征提取器
    
    用于mmWave雷达点云的空间结构特征提取
    """
    def __init__(self, input_dim: int = 3, hidden_dim: int = 64, k: int = 16):
        super().__init__()
        self.edge1 = EdgeConv(input_dim, hidden_dim, k)
        self.edge2 = EdgeConv(hidden_dim, hidden_dim * 2, k)
        self.edge3 = EdgeConv(hidden_dim * 2, hidden_dim * 4, k)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 点云序列, shape=(B, N, 3)
        Returns:
            features: 空间特征, shape=(B, N, 256)
        """
        x1 = self.edge1(x)      # (B, N, 64)
        x2 = self.edge2(x1)      # (B, N, 128)
        x3 = self.edge3(x2)       # (B, N, 256)
        return x3

class TransformerEncoder(nn.Module):
    """
    Transformer Encoder for temporal modeling of radar point cloud sequences
    
    将DGCNN提取的空间特征按时间维度进行全局建模
    """
    def __init__(
        self, 
        d_model: int = 256, 
        nhead: int = 8, 
        num_layers: int = 4,
        dim_feedforward: int = 1024,
        dropout: float = 0.1
    ):
        super().__init__()
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True,
            activation='gelu'
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)
        self.pos_encoder = PositionalEncoding(d_model, dropout)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 空间特征序列, shape=(B, T, N, d_model)
               T=时间步数, N=点数
        Returns:
            out: 时序特征, shape=(B, T, d_model)
        """
        B, T, N, D = x.shape
        
        # 对每个时间步的点特征做max pooling得到帧级特征
        frame_features = x.max(dim=2)[0]  # (B, T, D)
        
        # 加入位置编码
        frame_features = self.pos_encoder(frame_features)
        
        # Transformer编码
        out = self.encoder(frame_features)  # (B, T, D)
        
        # 时序聚合：取最后一帧 + 全局平均
        final_feature = out[:, -1, :] + out.mean(dim=1)  # (B, D)
        
        return final_feature


class PositionalEncoding(nn.Module):
    """正弦位置编码"""
    def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 500):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * 
            (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

class RadarChildPetClassifier(nn.Module):
    """
    完整的mmWave雷达儿童/宠物分类模型
    
    Architecture: DGCNN (空间) → Transformer (时序) → 分类头
    
    针对密闭空间（车内）设计的隐私保护CPD方案
    """
    def __init__(
        self,
        num_classes: int = 3,  # child, pet, empty
        input_dim: int = 3,    # x, y, z
        dgcnn_hidden: int = 64,
        transformer_dim: int = 256,
        nhead: int = 8,
        num_transformer_layers: int = 4,
        max_points: int = 1100,
        num_frames: int = 30,
        k: int = 16
    ):
        super().__init__()
        self.num_classes = num_classes
        self.max_points = max_points
        
        # 空间特征提取
        self.dgcnn = DGCNNFeatureExtractor(
            input_dim=input_dim,
            hidden_dim=dgcnn_hidden,
            k=k
        )
        
        # 时序建模
        self.transformer = TransformerEncoder(
            d_model=transformer_dim,
            nhead=nhead,
            num_layers=num_transformer_layers,
        )
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(transformer_dim, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, num_classes)
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 雷达点云序列, shape=(B, T, N, 3)
               B=batch, T=时间帧数, N=每帧点数
        Returns:
            logits: 分类输出, shape=(B, num_classes)
        """
        B, T, N, D = x.shape
        
        # 展平batch和时间为一个维度
        x_flat = x.reshape(B * T, N, D)
        
        # DGCNN空间特征
        spatial_features = self.dgcnn(x_flat)  # (B*T, N, 256)
        spatial_features = spatial_features.reshape(B, T, N, -1)
        
        # Transformer时序特征
        temporal_features = self.transformer(spatial_features)  # (B, 256)
        
        # 分类
        logits = self.classifier(temporal_features)
        return logits


# 实际测试
if __name__ == "__main__":
    import math
    
    # 模拟mmWave雷达数据
    model = RadarChildPetClassifier(
        num_classes=3,
        max_points=1100,
        num_frames=30
    )
    
    # 模拟输入: 30帧, 每帧1100个点, xyz坐标
    batch = torch.randn(4, 30, 1100, 3)
    
    # 前向传播
    output = model(batch)
    print(f"输入: {batch.shape}")
    print(f"输出: {output.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")
    
    # 推理速度测试
    import time
    model.eval()
    with torch.no_grad():
        start = time.time()
        for _ in range(100):
            _ = model(batch[:1])
        avg_time = (time.time() - start) / 100 * 1000
        print(f"推理延迟: {avg_time:.1f}ms (单帧30步序列)")

import numpy as np
from sklearn.cluster import DBSCAN

def preprocess_radar_pointcloud(
    frames: np.ndarray,
    eps: float = 0.4,
    min_samples: int = 6,
    z_weight: float = 0.25
) -> list:
    """
    mmWave雷达点云预处理管道
    
    Args:
        frames: 原始雷达帧列表, 每帧 shape=(N, 3)
        eps: DBSCAN邻域半径
        min_samples: DBSCAN最小点数
        z_weight: z轴权重（降低z轴影响，聚焦水平运动）
    
    Returns:
        processed_frames: 处理后的点云帧列表
    """
    processed = []
    
    for frame in frames:
        # Step 1: 去除零填充点
        norms = np.linalg.norm(frame, axis=1)
        valid_mask = norms ** 2 > 0.001
        points = frame[valid_mask]
        
        if len(points) < min_samples:
            processed.append(np.zeros((1, 3)))
            continue
        
        # Step 2: z轴加权（降低z轴影响）
        weighted_points = points.copy()
        weighted_points[:, 2] *= z_weight
        
        # Step 3: DBSCAN聚类
        clustering = DBSCAN(eps=eps, min_samples=min_samples).fit(weighted_points)
        labels = clustering.labels_
        
        # 保留最大簇
        unique_labels, counts = np.unique(labels[labels >= 0], return_counts=True)
        if len(unique_labels) > 0:
            largest_cluster = unique_labels[np.argmax(counts)]
            points = points[labels == largest_cluster]
        
        processed.append(points)
    
    return processed


# 测试
if __name__ == "__main__":
    # 模拟30帧雷达数据，每帧最多1100个点
    np.random.seed(42)
    frames = [np.random.randn(1100, 3) * 0.5 for _ in range(30)]
    # 加入一些零填充
    for f in frames:
        f[800:] = 0
    
    processed = preprocess_radar_pointcloud(frames)
    print(f"原始帧: {len(frames)}")
    print(f"处理后有效点数: {[len(p) for p in processed[:5]]}")