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

class AcousticMultiScaleEncoder(nn.Module):
    """
    多尺度声学编码器
    
    论文Section 3.2实现
    使用不同窗口大小的STFT提取互补的时频特征
    
    Args:
        signal_length: 输入信号长度
        n_persons: 最大人数
        n_joints: 关芽数（默认17）
    """
    
    def __init__(self, signal_length=16000, n_persons=3, n_joints=17):
        super().__init__()
        # 短窗口STFT - 捕获高时间分辨率
        # 窗口大小: 512 samples (~32ms @16kHz)
        self.short_window = 512
        self.short_hop = 160  # 10ms hop
        
        # 长窗口STFT - 捕获高频率分辨率
        # 窗口大小: 2048 samples (~128ms @16kHz)
        self.long_window = 2048
        self.long_hop = 320  # 20ms hop
        
        # 特征维度
        self.short_freq_bins = self.short_window // 2 + 1  # 257
        self.long_freq_bins = self.long_window // 2 + 1   # 1025
        
        # CNN骨干网络（类似ResNet）
        self.short_cnn = self._make_cnn_encoder(self.short_freq_bins)
        self.long_cnn = self._make_cnn_encoder(self.long_freq_bins)
        
        # 融合层
        self.fusion = nn.Sequential(
            nn.Conv2d(256 * 2, 256, kernel_size=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )
        
    def _make_cnn_encoder(self, freq_bins):
        """构建CNN编码器"""
        return nn.Sequential(
            # Input: (1, freq_bins, T_frames)
            nn.Conv2d(1, 64, kernel_size=(3, 3), padding=(1, 1)),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d((2, 2)),
            
            nn.Conv2d(64, 128, kernel_size=(3, 3), padding=(1, 1)),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d((2, 2)),
            
            nn.Conv2d(128, 256, kernel_size=(3, 3), padding=(1, 1)),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d((2, 2)),
            
            nn.Conv2d(256, 256, kernel_size=(3, 3), padding=(1, 1)),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
        )
    
    def compute_stft(self, x, window_size, hop_size):
        """计算STFT"""
        window = torch.hann_window(window_size, device=x.device)
        stft = torch.stft(
            x, 
            n_fft=window_size,
            hop_length=hop_size,
            win_length=window_size,
            window=window,
            return_complex=True
        )
        # 取幅度谱
        magnitude = torch.abs(stft)
        # log幅度
        magnitude = torch.log1p(magnitude)
        return magnitude.unsqueeze(1)  # (B, 1, F, T)
    
    def forward(self, x):
        """
        Args:
            x: 声学信号 (B, signal_length)
        Returns:
            encoded_features: (B, 256, T, 1)
        """
        # 多分辨率STFT
        short_spec = self.compute_stft(x, self.short_window, self.short_hop)
        long_spec = self.compute_stft(x, self.long_window, self.long_hop)
        
        # 对齐时间维度
        min_t = min(short_spec.size(-1), long_spec.size(-1))
        short_spec = short_spec[..., :min_t]
        long_spec = long_spec[..., :min_t]
        
        # CNN特征提取
        short_feat = self.short_cnn(short_spec)  # (B, 256, F', T)
        long_feat = self.long_cnn(long_spec)
        
        # 自适应池化到相同空间维度
        short_feat = F.adaptive_avg_pool2d(short_feat, (8, min_t))
        long_feat = F.adaptive_avg_pool2d(long_feat, (8, min_t))
        
        # 融合
        combined = torch.cat([short_feat, long_feat], dim=1)
        fused = self.fusion(combined)
        
        return fused  # (B, 256, 8, T)


class TemporalPoseDecoder(nn.Module):
    """
    时序姿态解码器
    
    论文Section 3.3实现
    使用注意力机制分离多人信息
    
    关键设计：
    - 每人每帧分配独立query
    - 时空注意力建模人体内时序动态
    - 交叉注意力建模人体间依赖
    """
    
    def __init__(self, d_model=256, n_persons=3, n_joints=17, 
                 n_heads=8, n_layers=3, seq_len=5):
        super().__init__()
        self.n_persons = n_persons
        self.n_joints = n_joints
        self.seq_len = seq_len  # 时序窗口帧数
        self.d_model = d_model
        
        # 每人每帧的learnable query
        # Shape: (n_persons * seq_len, d_model)
        self.person_queries = nn.Parameter(
            torch.randn(n_persons * seq_len, d_model) * 0.02
        )
        
        # 位置编码
        self.pos_encoding = nn.Parameter(
            torch.randn(1, n_persons * seq_len, d_model) * 0.02
        )
        
        # Transformer解码层
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=d_model * 4,
            dropout=0.1,
            batch_first=True
        )
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=n_layers)
        
        # 输出头: 3D关节坐标
        self.pose_head = nn.Linear(d_model, n_joints * 3)
        
        # 置信度预测
        self.confidence_head = nn.Linear(d_model, n_joints)
        
    def forward(self, memory):
        """
        Args:
            memory: 编码器输出 (B, S, d_model) 
                    S = spatial_dim * temporal_dim
        Returns:
            poses: (B, n_persons, n_joints, 3)
            confidences: (B, n_persons, n_joints)
        """
        B = memory.size(0)
        
        # 扩展query到batch维度
        queries = self.person_queries.unsqueeze(0).expand(B, -1, -1)
        queries = queries + self.pos_encoding
        
        # Transformer解码
        decoded = self.decoder(
            tgt=queries,       # (B, P*T, d_model)
            memory=memory      # (B, S, d_model)
        )  # (B, P*T, d_model)
        
        # Reshape: (B, P*T, d_model) -> (B, P, T, d_model)
        decoded = decoded.view(B, self.n_persons, self.seq_len, -1)
        
        # 取中间帧的输出作为当前帧姿态
        current = decoded[:, :, self.seq_len // 2, :]  # (B, P, d_model)
        
        # 预测3D姿态
        poses = self.pose_head(current)  # (B, P, n_joints*3)
        poses = poses.view(B, self.n_persons, self.n_joints, 3)
        
        # 预测置信度
        confidences = torch.sigmoid(
            self.confidence_head(current)
        )  # (B, P, n_joints)
        
        return poses, confidences


class SoundMHPE(nn.Module):
    """SoundMHPE: 完整的声学多人姿态估计模型"""
    
    def __init__(self, signal_length=16000, n_persons=3, 
                 n_joints=17, seq_len=5):
        super().__init__()
        self.encoder = AcousticMultiScaleEncoder(signal_length)
        self.decoder = TemporalPoseDecoder(
            d_model=256,
            n_persons=n_persons,
            n_joints=n_joints,
            seq_len=seq_len
        )
        
    def forward(self, x):
        """
        Args:
            x: 声学信号 (B, signal_length)
        Returns:
            poses: (B, n_persons, n_joints, 3)
            confidences: (B, n_persons, n_joints)
        """
        # 编码
        features = self.encoder(x)  # (B, 256, 8, T)
        B, C, H, T = features.shape
        
        # 展平为序列
        memory = features.permute(0, 3, 1, 2)  # (B, T, C, H)
        memory = memory.reshape(B, T, C * H)   # (B, T, 2048)
        
        # 如果维度不匹配，添加投影
        if memory.size(-1) != 256:
            memory = nn.Linear(memory.size(-1), 256).to(memory.device)(memory)
        
        # 解码
        poses, confidences = self.decoder(memory)
        
        return poses, confidences


# 测试代码
if __name__ == "__main__":
    # 模拟参数
    batch_size = 4
    signal_length = 16000  # 1秒 @ 16kHz
    n_persons = 3
    n_joints = 17
    
    # 创建模型
    model = SoundMHPE(
        signal_length=signal_length,
        n_persons=n_persons,
        n_joints=n_joints
    )
    
    # 模拟声学输入
    x = torch.randn(batch_size, signal_length)
    
    # 前向传播
    poses, confs = model(x)
    
    print(f"输入信号: {x.shape}")
    print(f"输出姿态: {poses.shape}")  # (4, 3, 17, 3)
    print(f"置信度: {confs.shape}")   # (4, 3, 17)
    print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M")
    
    # 计算MPJPE
    # 假设ground truth
    gt_poses = torch.randn(batch_size, n_persons, n_joints, 3)
    mpjpe = torch.mean(torch.norm(poses - gt_poses, dim=-1))  # mm
    print(f"模拟MPJPE: {mpjpe:.1f} mm (论文报告: 106.5 mm)")

# 量化压缩方案
import torch.quantization as quant

def optimize_for_edge(model):
    """边缘部署优化"""
    # 1. 动态量化（INT8）
    quantized = quant.quantize_dynamic(
        model,
        {nn.Linear, nn.Conv2d},
        dtype=quant.qint8
    )
    
    # 2. 模型大小对比
    original = sum(p.numel() for p in model.parameters()) * 4  # float32
    quantized_size = sum(
        p.qsize() if hasattr(p, 'qsize') else p.numel() * 4
        for p in quantized.parameters()
    )
    
    print(f"原始模型: {original / 1e6:.2f} MB")
    print(f"量化模型: {quantized_size / 1e6:.2f} MB")
    print(f"压缩比: {original / quantized_size:.1f}x")
    
    return quantized

# 3. 实时推理配置
INFERENCE_CONFIG = {
    "signal_duration": 0.5,  # 0.5秒信号
    "sample_rate": 16000,    # 16kHz
    "buffer_size": 8000,     # 0.5s buffer
    "inference_interval": 1.0,  # 每秒推理一次
    "max_persons": 3,        # 最多3人
    "joints": 17,            # 17关节
    "device": "npu",         # NPU推理
    "precision": "int8",     # INT8量化
}