"""Optional pooling heads for the research/ablation track.

DEPLOYMENT NOTE
---------------
The production model MUST remain a stock `DebertaV2ForSequenceClassification`
so the HF Inference API can serve it through the standard `text-classification`
pipeline (no `trust_remote_code` on the Inference API). Everything in this file
is for the ablation table in the paper, not for the deployed checkpoint.

Motivation: DeBERTa's default `ContextPooler` reads only the [CLS] token. That is
reasonable for a tweet, where the emotional payload is one clause. A verse of
lyrics spreads the signal across the whole passage — "and at once I knew / I was
not magnificent" carries its weight in the second line, not the first token. Mean
and attention pooling recover that.

Use `--pooling attention` in train_v2.py to run the ablation.
"""

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoConfig, AutoModel
from transformers.modeling_outputs import SequenceClassifierOutput


class AttentionPooling(nn.Module):
    """Learned additive attention over token hidden states.

    Computes a scalar score per token, softmaxes over the (masked) sequence, and
    returns the weighted sum. Unlike [CLS] pooling this lets the model learn that
    e.g. "hollow", "still", "again" carry the emotional load in a lyric line.
    """

    def __init__(self, hidden_size: int, dropout: float = 0.1):
        super().__init__()
        self.attn = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.Tanh(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size, 1),
        )

    def forward(self, hidden: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        # hidden: (B, T, H), mask: (B, T)
        scores = self.attn(hidden).squeeze(-1)               # (B, T)
        scores = scores.masked_fill(mask == 0, torch.finfo(scores.dtype).min)
        weights = F.softmax(scores, dim=-1).unsqueeze(-1)    # (B, T, 1)
        return (hidden * weights).sum(dim=1)                 # (B, H)


class MeanMaxAttentionPool(nn.Module):
    """Concatenation of masked-mean, masked-max, and attention pooling.

    The three views are complementary: mean captures overall tone, max captures
    the single most charged token, attention learns a task-specific weighting.
    Output dimension is 3 * hidden_size.
    """

    def __init__(self, hidden_size: int, dropout: float = 0.1):
        super().__init__()
        self.attention = AttentionPooling(hidden_size, dropout)
        self.output_size = hidden_size * 3

    def forward(self, hidden: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        mask_f = mask.unsqueeze(-1).float()

        summed = (hidden * mask_f).sum(dim=1)
        counts = mask_f.sum(dim=1).clamp(min=1e-9)
        mean_pooled = summed / counts

        max_pooled = hidden.masked_fill(mask.unsqueeze(-1) == 0, torch.finfo(hidden.dtype).min)
        max_pooled = max_pooled.max(dim=1).values

        attn_pooled = self.attention(hidden, mask)
        return torch.cat([mean_pooled, max_pooled, attn_pooled], dim=-1)


class MultiSampleDropoutHead(nn.Module):
    """Classifier that averages logits over several independent dropout masks.

    Cheap variance reduction: one forward pass through the backbone, N cheap
    passes through the linear head. Converges faster and generalizes better than
    a single-mask head, which matters when the tail classes have few examples.

    Reference: Inoue 2019, "Multi-Sample Dropout for Accelerated Training and
    Better Generalization" (arXiv:1905.09788)
    """

    def __init__(self, input_size: int, num_labels: int, num_samples: int = 5, p: float = 0.3):
        super().__init__()
        self.dropouts = nn.ModuleList([nn.Dropout(p) for _ in range(num_samples)])
        self.classifier = nn.Linear(input_size, num_labels)

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        return torch.stack(
            [self.classifier(d(features)) for d in self.dropouts], dim=0
        ).mean(dim=0)


class EumoraClassifier(nn.Module):
    """DeBERTa backbone + rich pooling + multi-sample dropout head.

    Research-track model. Same backbone weights as the deployed checkpoint, so a
    TAPT-adapted encoder can be dropped into either path for a clean ablation.
    """

    def __init__(
        self,
        model_name: str,
        num_labels: int,
        pooling: str = "attention",
        dropout: float = 0.3,
        num_dropout_samples: int = 5,
    ):
        super().__init__()
        self.config = AutoConfig.from_pretrained(model_name, num_labels=num_labels)
        self.backbone = AutoModel.from_pretrained(model_name, config=self.config)
        hidden = self.config.hidden_size
        self.pooling_type = pooling

        if pooling == "attention":
            self.pool = MeanMaxAttentionPool(hidden, dropout=0.1)
            head_input = self.pool.output_size
        elif pooling == "mean":
            self.pool = None
            head_input = hidden
        else:  # "cls"
            self.pool = None
            head_input = hidden

        self.head = MultiSampleDropoutHead(
            head_input, num_labels, num_samples=num_dropout_samples, p=dropout
        )

    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
        hidden = self.backbone(
            input_ids=input_ids, attention_mask=attention_mask
        ).last_hidden_state

        if self.pooling_type == "attention":
            features = self.pool(hidden, attention_mask)
        elif self.pooling_type == "mean":
            mask_f = attention_mask.unsqueeze(-1).float()
            features = (hidden * mask_f).sum(dim=1) / mask_f.sum(dim=1).clamp(min=1e-9)
        else:
            features = hidden[:, 0]

        logits = self.head(features)
        return SequenceClassifierOutput(logits=logits)