"""Mask-aware temporal pooling shared by student and teacher models."""

from __future__ import annotations

import torch
from torch import Tensor, nn


class MaskedAttentiveStatisticsPooling(nn.Module):
    """Learned weighted mean and standard deviation over valid frames.

    Input is ``[batch, channels, frames]`` and output is
    ``[batch, 2 * channels]``.  Multiplicative masking keeps even an all-padded
    item finite, which is useful when robustly serving malformed audio.
    """

    def __init__(self, channels: int, attention_channels: int = 128) -> None:
        super().__init__()
        self.attention = nn.Sequential(
            nn.Conv1d(channels, attention_channels, kernel_size=1),
            nn.Tanh(),
            nn.Conv1d(attention_channels, 1, kernel_size=1),
        )

    def forward(self, x: Tensor, mask: Tensor | None = None) -> Tensor:
        if x.ndim != 3:
            raise ValueError("pooling input must have shape [batch, channels, frames]")
        if mask is None:
            mask = torch.ones((x.shape[0], x.shape[-1]), dtype=torch.bool, device=x.device)
        if mask.shape != (x.shape[0], x.shape[-1]):
            raise ValueError("mask must have shape [batch, frames]")

        valid = mask.unsqueeze(1).to(x.dtype)
        scores = self.attention(x)
        scores = scores.masked_fill(~mask.unsqueeze(1), -1e4)
        scores = scores - scores.amax(dim=-1, keepdim=True)
        weights = torch.exp(scores) * valid
        weights = weights / weights.sum(dim=-1, keepdim=True).clamp_min(1e-8)

        mean = (weights * x).sum(dim=-1)
        second_moment = (weights * x.square()).sum(dim=-1)
        std = (second_moment - mean.square()).clamp_min(1e-5).sqrt()
        return torch.cat([mean, std], dim=1)