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

class DualStreamGestureNet(nn.Module):
    """
    双流 LSTM + Attention-BiGRU 手势识别网络
    
    Stream 1: LSTM 处理距离-多普勒时序
    Stream 2: BiGRU + Attention 处理角度时序
    """
    
    def __init__(self, input_size_range=(64, 64), 
                 input_size_angle=(32,), 
                 hidden_dim=128, num_classes=8):
        super().__init__()
        
        # Stream 1: LSTM for Range-Doppler
        self.lstm = nn.LSTM(
            input_size=input_size_range[0] * input_size_range[1],
            hidden_size=hidden_dim,
            num_layers=2,
            batch_first=True,
            bidirectional=False,
        )
        
        # Stream 2: BiGRU + Attention for Angle
        self.bigru = nn.GRU(
            input_size=input_size_angle[0],
            hidden_size=hidden_dim,
            num_layers=2,
            batch_first=True,
            bidirectional=True,
        )
        
        # Attention mechanism
        self.attention = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, 1),
            nn.Softmax(dim=1),
        )
        
        # Fusion + Classification
        self.fusion = nn.Sequential(
            nn.Linear(hidden_dim + hidden_dim * 2, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, num_classes),
        )
    
    def forward(self, range_doppler, angle_seq):
        """
        Args:
            range_doppler: (B, T, H, W) 距离-多普勒时序
            angle_seq: (B, T, A) 角度时序
        """
        B, T, H, W = range_doppler.shape
        
        # Stream 1: LSTM
        rd_flat = range_doppler.view(B, T, H * W)
        lstm_out, _ = self.lstm(rd_flat)
        lstm_feat = lstm_out[:, -1, :]  # 取最后时步
        
        # Stream 2: BiGRU + Attention
        gru_out, _ = self.bigru(angle_seq)
        
        # Attention weighting
        att_weights = self.attention(gru_out)  # (B, T, 1)
        attended = torch.sum(att_weights * gru_out, dim=1)  # (B, hidden*2)
        
        # Fusion
        combined = torch.cat([lstm_feat, attended], dim=1)
        output = self.fusion(combined)
        
        return output, att_weights


# 测试
if __name__ == "__main__":
    model = DualStreamGestureNet(num_classes=8)
    
    # 模拟输入
    B, T = 4, 30  # 批次4, 时序30帧
    range_doppler = torch.randn(B, T, 64, 64)
    angle_seq = torch.randn(B, T, 32)
    
    output, att_weights = model(range_doppler, angle_seq)
    
    print(f"输入: RD={range_doppler.shape}, Angle={angle_seq.shape}")
    print(f"输出: {output.shape}")
    print(f"注意力权重: {att_weights.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")