import torch
import torch.nn as nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence


class BiLSTMFeatureExtractor(nn.Module):
    """
    BiLSTM Feature Extractor module.
    
    Inputs:
        x    : Padded WavLM sequence embeddings [B, T, 768]
        mask : Timestep mask [B, T] (optional, 1 for real timesteps, 0 for padding)
        
    Outputs:
        out  : BiLSTM contextual representations [B, T, 256]
    """
    def __init__(self, input_size=768, hidden_size=128, num_layers=1, dropout=0.0):
        super(BiLSTMFeatureExtractor, self).__init__()
        
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bidirectional = True
        
        # BiLSTM Layer: hidden_size=128, bidirectional=True -> Output dim = 128 * 2 = 256
        self.bilstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=True,
            dropout=dropout if num_layers > 1 else 0.0
        )
        
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape
        
        if mask is not None:
            lengths = mask.sum(dim=1).cpu()
            
            packed_x = pack_padded_sequence(
                x,
                lengths,
                batch_first=True,
                enforce_sorted=False
            )
            
            packed_out, (hn, cn) = self.bilstm(packed_x)
            
            out, _ = pad_packed_sequence(
                packed_out,
                batch_first=True,
                total_length=seq_len
            )
        else:
            out, (hn, cn) = self.bilstm(x)
            
        return out


class TemporalAttention(nn.Module):
    """
    Temporal Attention Module.
    Learns frame-wise attention scores over sequence length T to aggregate
    BiLSTM output sequence [B, T, 256] into a fixed 256-dimensional context vector [B, 256].
    
    Inputs:
        h    : BiLSTM outputs [B, T, input_dim] (input_dim=256)
        mask : Timestep mask [B, T] (1 for real timesteps, 0 for padding)
        
    Outputs:
        context      : Aggregated context vector [B, 256]
        attn_weights : Normalized attention weights [B, T]
    """
    def __init__(self, input_dim=256):
        super(TemporalAttention, self).__init__()
        self.input_dim = input_dim
        self.w = nn.Linear(input_dim, 1, bias=False)
        
    def forward(self, h, mask=None):
        scores = self.w(torch.tanh(h)).squeeze(-1)  # [B, T]
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
            
        attn_weights = torch.softmax(scores, dim=1)  # [B, T]
        context = torch.bmm(attn_weights.unsqueeze(1), h).squeeze(1)  # [B, 256]
        
        return context, attn_weights


class BiLSTMAttentionClassifier(nn.Module):
    """
    Full Downstream Model Architecture:
    WavLM Embedding [B, T, 768] -> BiLSTM -> [B, T, 256] -> Temporal Attention -> [B, 256] -> Linear Classifier -> [B, 6] Logits
    """
    def __init__(self, input_size=768, hidden_size=128, num_classes=6, dropout=0.3):
        super(BiLSTMAttentionClassifier, self).__init__()
        
        self.bilstm = BiLSTMFeatureExtractor(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=1,
            dropout=dropout
        )
        
        context_dim = hidden_size * 2  # 256
        self.attention = TemporalAttention(input_dim=context_dim)
        
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(context_dim, num_classes)
        
    def forward(self, x, mask=None):
        # 1. BiLSTM: [B, T, 768] -> [B, T, 256]
        bilstm_out = self.bilstm(x, mask=mask)
        
        # 2. Temporal Attention: [B, T, 256] -> [B, 256]
        context, attn_weights = self.attention(bilstm_out, mask=mask)
        
        # 3. Dropout + Linear Classifier: [B, 256] -> [B, 6]
        dropped_context = self.dropout(context)
        logits = self.classifier(dropped_context)
        
        return logits, attn_weights