# models/meralion_encoder.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Optional

from transformers import AutoModel,PretrainedConfig, PreTrainedModel, AutoConfig
from peft import get_peft_model, LoraConfig
from omegaconf import DictConfig

# 1. Define a Config class that holds all your YAML settings
class MeralionGenderConfig(PretrainedConfig):
    model_type = "meralion_gender"
    def __init__(
        self, 
        meralion_name="", 
        num_classes=2, 
        aggregator="attention", 
        downstream_params=None,
        **kwargs
    ):
        # Pass all basic types (strings, ints, dicts) to super
        super().__init__(
            meralion_name=meralion_name,
            num_classes=num_classes,
            aggregator=aggregator,
            downstream_params=downstream_params or {},
            **kwargs
        )

class SE1d(nn.Module):
    """Squeeze-and-Excitation block for 1D convolutions"""
    def __init__(self, channels: int, reduction: int = 8):
        super().__init__()
        hidden = max(8, channels // reduction)
        self.avg = nn.AdaptiveAvgPool1d(1)
        self.fc = nn.Sequential(
            nn.Conv1d(channels, hidden, 1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv1d(hidden, channels, 1, bias=False),
            nn.Sigmoid(),
        )
    
    def forward(self, x: torch.Tensor):
        w = self.fc(self.avg(x))
        return x * w


class Res2Block1d(nn.Module):
    """Res2Net block adapted for 1D convolutions - BatchNorm free"""
    def __init__(self, channels: int, scale: int = 4, kernel_size: int = 3, dilation: int = 1):
        super().__init__()
        assert channels % scale == 0, f"channels ({channels}) must be divisible by scale ({scale})"
        self.scale = scale
        self.width = channels // scale
        pad = (kernel_size // 2) * dilation
        
        self.convs = nn.ModuleList([
            nn.Conv1d(self.width, self.width, kernel_size, padding=pad, dilation=dilation, bias=True)
            for _ in range(scale - 1)
        ])
        self.norm = nn.GroupNorm(num_groups=min(32, channels), num_channels=channels)
        self.act = nn.ReLU(inplace=True)
    
    def forward(self, x: torch.Tensor):
        xs = torch.split(x, self.width, dim=1)
        out = [xs[0]]
        for i, conv in enumerate(self.convs, start=1):
            if i == 1:
                s = xs[i]
            else:
                s = xs[i] + out[-1]  # Fixed: proper residual connection
            out.append(conv(s))
        y = torch.cat(out, dim=1)
        return self.act(self.norm(y))

class ECAPABlock(nn.Module):
    """Enhanced ECAPA block with proper residual connections - BatchNorm free"""
    def __init__(self, channels: int, scale: int = 4, kernel_size: int = 3, dilation: int = 1):
        super().__init__()
        self.conv1 = nn.Conv1d(channels, channels, 1, bias=True)
        self.norm1 = nn.GroupNorm(num_groups=min(32, channels), num_channels=channels)
        self.act1 = nn.ReLU(inplace=True)
        
        self.res2 = Res2Block1d(channels, scale=scale, kernel_size=kernel_size, dilation=dilation)
        self.se = SE1d(channels)
        
        self.conv2 = nn.Conv1d(channels, channels, 1, bias=True)
        self.norm2 = nn.GroupNorm(num_groups=min(32, channels), num_channels=channels)
        self.act2 = nn.ReLU(inplace=True)
    
    def forward(self, x: torch.Tensor):
        residual = x
        
        y = self.act1(self.norm1(self.conv1(x)))
        y = self.res2(y)
        y = self.se(y)
        y = self.norm2(self.conv2(y))
        
        return self.act2(y + residual)


class EmotionECAPATDNN(nn.Module):
    """ECAPA-TDNN optimized for emotion recognition with hierarchical attention"""
    def __init__(
        self, 
        input_dim: int, 
        channels: int = 512,
        output_dim: int = 256,
        num_blocks: int = 3,
        dilations: tuple = (1, 2, 3),
        embed_dim: int = 512,
        num_emotions: int = 8,  # Common emotion categories
        dropout: float = 0.2,
        pooling_type="attention"
    ):
        super().__init__()
        self.output_dim = output_dim
        # Input projection with layer norm for stability
        self.proj_in = nn.Sequential(
            nn.Linear(input_dim, channels),
            nn.LayerNorm(channels),
            nn.GELU(),
            nn.Dropout(0.2),
        )
        
        # ECAPA blocks with different dilations
        self.blocks = nn.ModuleList([
            ECAPABlock(channels, scale=4, kernel_size=3, dilation=d) 
            for d in dilations
        ])
        
        # Multi-scale feature aggregation
        self.mfa = nn.Sequential(
            nn.Conv1d(channels * (len(dilations) + 1), channels, 1, bias=True),
            nn.GroupNorm(num_groups=min(32, channels), num_channels=channels),
            nn.GELU()
            #nn.ReLU(inplace=True)
        )

        self.pooling = AttentionPooling(channels)
        
        # Final embedding layers
        self.embed = nn.Sequential(
            nn.Linear(channels, channels),
            nn.LayerNorm(channels),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(channels, channels // 2),
            nn.LayerNorm(channels // 2),
            nn.GELU(),
            nn.Dropout(dropout)
        )
        
        # Initialize weights
        self._init_weights()
    
    def _init_weights(self):
        """Initialize model weights"""
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Conv1d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
    
    def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, return_embeddings: bool = False):
        """
        Forward pass
        
        Args:
            x: Input tensor (batch, time, features) - from Whisper encoder
            attention_mask: Attention mask (batch, time)
            return_embeddings: Whether to return embeddings instead of logits
        
        Returns:
            If return_embeddings=False: emotion logits (batch, num_emotions)
            If return_embeddings=True: feature embeddings (batch, embed_dim // 2)
        """
        # Project input and transpose for conv1d
        x = self.proj_in(x)  # (batch, time, channels)
        x_conv = x.transpose(1, 2)  # (batch, channels, time)
        
        # Apply ECAPA blocks and collect multi-scale features
        features = [x_conv]
        for block in self.blocks:
            x_conv = block(x_conv)
            features.append(x_conv)
        
        # Multi-scale feature aggregation
        y = torch.cat(features, dim=1)  # (batch, channels * (n_blocks + 1), time)
        y = self.mfa(y)  # (batch, channels, time)
        y = y.transpose(1, 2)  # (batch, time, channels)
        
        # Hierarchical attention pooling
        #pooled = self.norm_layer(self.pooling(y))  # (batch, channels)
        pooled = self.pooling(y)  # (batch, channels)

        # Generate embeddings
        embeddings = self.embed(pooled)  # (batch, embed_dim // 2)
        
        return embeddings
        

class LayerAttentiveAggregation(nn.Module):
    """
    Smart Layer Aggregation:
    Instead of a static weighted sum, this computes attention weights 
    based on the hidden states themselves.
    """
    def __init__(self, hidden_size: int, num_layers: int):
        super().__init__()
        # Transformation to compute score per layer
        self.query_proj = nn.Linear(hidden_size, 1)
        self.num_layers = num_layers

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # hidden_states: (L, B, T, D)
        # We want to learn which layer 'L' is most important per timestep or globally.
        # Let's do global (per sample) importance to save compute.
        
        # Mean pool over time for the scoring mechanism: (L, B, D)
        # Using mean helps avoid noise from silence frames
        global_repr = hidden_states.mean(dim=2) 
        
        # Compute scores: (L, B, 1)
        scores = self.query_proj(global_repr)
        
        # Softmax over layers (dim=0) -> (L, B, 1)
        attn_weights = F.softmax(scores, dim=0)
        
        # Reshape for broadcasting: (L, B, 1, 1)
        attn_weights = attn_weights.unsqueeze(-1)
        
        # Weighted sum: sum((L, B, T, D) * (L, B, 1, 1)) -> (B, T, D)
        aggregated = (hidden_states * attn_weights).sum(dim=0)
        #print(f"DEBUG: Number of hidden states provided to aggregator: {len(hidden_states)}")
        #exit()
        return aggregated

class MeralionForGenderClassification(PreTrainedModel):
    config_class = MeralionGenderConfig

    def __init__(self, config: MeralionGenderConfig):
        super().__init__(config)
        
        
        # 1. Load Backbone        
        backbone_config = AutoConfig.from_pretrained(config.meralion_name, trust_remote_code=True)
        self.backbone = AutoModel.from_config(backbone_config, trust_remote_code=True)
        
        hidden_size = getattr(self.backbone.config, "hidden_size", None) \
                   or getattr(self.backbone.config, "d_model", None)
        if hidden_size is None:
            raise ValueError("Cannot infer hidden size from MERaLiON config.")
        
        num_layers = getattr(self.backbone.config, "num_hidden_layers", 0) + 1 # +1 for embeddings

        # Reset all to requires_grad=True first
        for p in self.backbone.parameters():
            p.requires_grad = False

        # 3. Layer Aggregation
        self.backbone.config.output_hidden_states = True
        self.layer_aggregator = LayerAttentiveAggregation(hidden_size, num_layers)

        # 4. Downstream Head
        self.downstream = EmotionECAPATDNN(
            input_dim=hidden_size, 
        )
        d_out = self.downstream.output_dim

        # 5. Gender Heads
        self.gender_proj  = nn.Linear(d_out, 256)
        self.gender_head  = nn.Sequential(
            nn.RMSNorm(256), nn.GELU(), nn.Linear(256, config.num_classes)
        )

    def _forward_backbone(self, inputs: torch.Tensor, attention_mask: torch.Tensor):
        # During inference, we always want hidden states for the aggregator
        outputs = self.backbone(
            input_values=inputs, 
            attention_mask=attention_mask, 
            output_hidden_states=True
        )
        
        # Your Aggregation Logic
        hs = torch.stack(outputs.hidden_states, dim=0) # (L, B, T, D)
        return self.layer_aggregator(hs)

    def forward(
        self, 
        input_values: torch.Tensor, 
        attention_mask: torch.Tensor,
        **kwargs
    ):

        inputs = input_values 
        x = self._forward_backbone(inputs, attention_mask)   # (B, T, D)
        feats = self.downstream(x)                   # Do Not Pass mask to downstream
        
        pre_final = self.gender_proj(feats)
        logits    = self.gender_head(pre_final)
        return pre_final, logits
        
   
class AttentionPooling(nn.Module):
    """
    Attention-based pooling over the sequence dimension.
    Input: (batch, seq_len, embed_dim)
    Output: (batch, embed_dim)
    """
    def __init__(self, embed_dim):
        super().__init__()
        self.attention = nn.Linear(embed_dim, 1)

    def forward(self, x, mask=None):
        # x: (batch, seq_len, embed_dim)
        attn_scores = self.attention(x).squeeze(-1)  # (batch, seq_len)
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))
        attn_weights = torch.softmax(attn_scores, dim=1)  # (batch, seq_len)
        pooled = torch.sum(x * attn_weights.unsqueeze(-1), dim=1)  # (batch, embed_dim)
        return pooled