from __future__ import annotations

import torch
from torch import nn
from transformers import AutoConfig, AutoModel


class EncoderClassifier(nn.Module):
    """Small encoder + sentence pooling + a linear classification head."""

    def __init__(
        self,
        base_model: str,
        num_labels: int,
        dropout: float = 0.1,
        pooling: str = "mean",
        pretrained: bool = True,
        encoder_config: dict | None = None,
    ) -> None:
        super().__init__()
        if pretrained:
            self.encoder = AutoModel.from_pretrained(base_model)
        else:
            if encoder_config is None:
                raise ValueError("encoder_config is required when pretrained=False")
            raw_config = dict(encoder_config)
            model_type = raw_config.pop("model_type")
            config = AutoConfig.for_model(model_type, **raw_config)
            self.encoder = AutoModel.from_config(config)
        self.dropout = nn.Dropout(dropout)
        self.pooling = pooling
        self.classifier = nn.Linear(self.encoder.config.hidden_size, num_labels)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        hidden = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
        ).last_hidden_state
        if self.pooling == "cls":
            pooled = hidden[:, 0]
        elif self.pooling == "mean":
            mask = attention_mask.unsqueeze(-1).to(hidden.dtype)
            pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
        else:
            raise ValueError(f"Unsupported pooling: {self.pooling}")
        return self.classifier(self.dropout(pooled))