"""Mask-aware pooling modules."""

from __future__ import annotations

import torch
from torch import nn


def masked_mean(hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
    mask = attention_mask.to(hidden_states.dtype).unsqueeze(-1)
    denom = mask.sum(dim=1).clamp_min(1.0)
    return (hidden_states * mask).sum(dim=1) / denom


class MaskedAttentionPool(nn.Module):
    def __init__(self, hidden_size: int):
        super().__init__()
        self.query = nn.Parameter(torch.empty(hidden_size))
        self.proj = nn.Linear(hidden_size, hidden_size)
        nn.init.normal_(self.query, mean=0.0, std=0.02)

    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        scores = torch.einsum("blh,h->bl", torch.tanh(self.proj(hidden_states)), self.query)
        scores = scores.masked_fill(~attention_mask, -1.0e4)
        weights = torch.softmax(scores, dim=-1) * attention_mask.to(hidden_states.dtype)
        weights = weights / weights.sum(dim=-1, keepdim=True).clamp_min(1.0e-6)
        return torch.einsum("bl,blh->bh", weights, hidden_states)


class StrataBertPooler(nn.Module):
    def __init__(self, hidden_size: int, pooling_type: str):
        super().__init__()
        self.pooling_type = pooling_type
        self.attention_pool = MaskedAttentionPool(hidden_size)
        if pooling_type == "hybrid_masked":
            self.out = nn.Linear(hidden_size * 3, hidden_size)
        elif pooling_type in {"cls", "masked_mean", "masked_attention"}:
            self.out = nn.Identity()
        else:
            raise ValueError(f"unknown pooling_type: {pooling_type}")

    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        cls = hidden_states[:, 0]
        mean = masked_mean(hidden_states, attention_mask)
        attended = self.attention_pool(hidden_states, attention_mask)
        if self.pooling_type == "cls":
            return cls
        if self.pooling_type == "masked_mean":
            return mean
        if self.pooling_type == "masked_attention":
            return attended
        return self.out(torch.cat([cls, mean, attended], dim=-1))