"""
Shared NativityClassifier for training and inference.
Variants: MLP (baseline), MLP + channel attention (Squeeze-and-Excitation style).
"""
import torch
import torch.nn as nn


class ImprovedAttentionBlock(nn.Module):
    """
    Squeeze-and-Excitation style attention: squeeze 1024 dims to a bottleneck,
    then excite back to 1024 channel weights. Learns which channels matter for nativity.
    """
    def __init__(self, input_dim=1024, reduction_ratio=8):
        super().__init__()
        reduced = max(input_dim // reduction_ratio, 32)
        self.fc1 = nn.Linear(input_dim, reduced)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(reduced, input_dim)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # x shape: (batch_size, 1024)
        attention_weights = self.sigmoid(self.fc2(self.relu(self.fc1(x))))
        return x * attention_weights


class NativityClassifier(nn.Module):
    """
    MLP classifier on wav2vec embeddings -> Native vs Non-Native.
    Uses BatchNorm and dropout for better generalization.
    """

    def __init__(
        self,
        input_dim=1024,
        hidden_dims=(384, 256, 128),
        output_dim=2,
        dropout=0.35,
    ):
        super(NativityClassifier, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim

        layers = []
        prev = input_dim
        for h in hidden_dims:
            layers.extend([
                nn.Linear(prev, h),
                nn.BatchNorm1d(h),
                nn.ReLU(inplace=True),
                nn.Dropout(dropout),
            ])
            prev = h
        layers.append(nn.Linear(prev, output_dim))
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)


class NativityClassifierWithAttention(nn.Module):
    """
    Squeeze-and-Excitation attention over 1024-d embedding then MLP.
    SE block learns which channels matter for non-native prosody; higher final_dropout
    reduces overfitting to majority class.
    """

    def __init__(
        self,
        input_dim=1024,
        hidden_dims=(384, 256, 128),
        output_dim=2,
        dropout=0.35,
        final_dropout=0.6,
        reduction_ratio=8,
    ):
        super(NativityClassifierWithAttention, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.attention = ImprovedAttentionBlock(input_dim=input_dim, reduction_ratio=reduction_ratio)
        layers = []
        prev = input_dim
        for i, h in enumerate(hidden_dims):
            layers.extend([
                nn.Linear(prev, h),
                nn.BatchNorm1d(h),
                nn.ReLU(inplace=True),
                nn.Dropout(final_dropout if i == len(hidden_dims) - 1 else dropout),
            ])
            prev = h
        layers.append(nn.Linear(prev, output_dim))
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        x = self.attention(x)
        return self.net(x)