"""Shared classifier heads for BERT-family encoders.

These three head classes are backbone-agnostic - they accept any HuggingFace
``model_name`` and load the appropriate encoder via ``AutoModel.from_pretrained``.
The same head can be stacked on IndoBERT, XLM-R, mDeBERTa, IndoBERTweet, etc.
Used by the four hybrid-transformer families (one per encoder backbone):

- :mod:`src.models.indobertweet_cnn` / ``_bilstm`` / ``_cnn_bilstm`` (IBT)
- :mod:`src.models.indobert_hybrid`        (IndoBERT)
- :mod:`src.models.xlm_roberta_hybrid`     (XLM-R)
- :mod:`src.models.mdeberta_hybrid`        (mDeBERTa)

Each head sits between the encoder's ``last_hidden_state`` (B, T, H) and a
single-logit ``Linear`` projection for BCE binary classification.

References:
- Kusuma & Chowanda (2023, JOIV) - IBT+CNN and IBT+BiLSTM heads for
  Indonesian hate speech detection.
- Mozafari, Farahbakhsh & Crespi (2020) - BERT+CNN for hate speech.
- Faris, Aljarah & Castillo (2020) - BERT-CNN-BiLSTM stack for Arabic.
"""

from __future__ import annotations

import torch
import torch.nn as nn
from transformers import AutoModel


def _load_encoder(model_name: str) -> AutoModel:
    """Load a HuggingFace encoder in fp32.

    Some checkpoints (notably ``microsoft/mdeberta-v3-base``) ship safetensors
    in fp16; the classification head is fp32, and that mismatch crashes
    ``F.linear(pooled_fp16, weight_fp32)``. Cast end-to-end fp32.
    """
    bert = AutoModel.from_pretrained(model_name, use_safetensors=True)
    return bert.float()


class BertCNNHead(nn.Module):
    """Encoder -> Conv1d -> ReLU -> Dropout -> MaxPool -> Linear (single logit).

    Replicates the IndoBERTweet+CNN head from Kusuma & Chowanda (2023).
    Captures local n-gram patterns (kernel_size=3 → trigrams) via the
    Conv1d over the token-axis of the encoder's last hidden state.
    """

    def __init__(
        self,
        model_name: str,
        dropout_rate: float,
        cnn_filters: int = 128,
        cnn_kernel_size: int = 3,
    ) -> None:
        super().__init__()
        self.bert = _load_encoder(model_name)
        hidden = self.bert.config.hidden_size
        self.cnn = nn.Conv1d(
            in_channels=hidden,
            out_channels=cnn_filters,
            kernel_size=cnn_kernel_size,
            padding=cnn_kernel_size // 2,
        )
        self.dropout = nn.Dropout(dropout_rate)
        self.pool = nn.AdaptiveMaxPool1d(1)
        self.classifier = nn.Linear(cnn_filters, 1)

    def forward(self, input_ids: torch.Tensor,
                attention_mask: torch.Tensor) -> torch.Tensor:
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        h = outputs.last_hidden_state  # (B, T, hidden)
        c = self.cnn(h.transpose(1, 2))  # (B, cnn_filters, T)
        c = torch.relu(c)
        c = self.dropout(c)
        pooled = self.pool(c).squeeze(-1)  # (B, cnn_filters)
        return self.classifier(pooled)  # (B, 1)


class BertBiLSTMHead(nn.Module):
    """Encoder -> BiLSTM -> Dropout -> MaxPool -> Linear (single logit).

    Replicates the IndoBERTweet+BiLSTM head from Kusuma & Chowanda (2023),
    which was the paper's best variant (93.3% F1 on Alfina et al.).
    Captures bidirectional sequence dependencies after the transformer.
    """

    def __init__(
        self,
        model_name: str,
        dropout_rate: float,
        lstm_units: int = 64,
    ) -> None:
        super().__init__()
        self.bert = _load_encoder(model_name)
        hidden = self.bert.config.hidden_size
        self.bilstm = nn.LSTM(
            input_size=hidden,
            hidden_size=lstm_units,
            batch_first=True,
            bidirectional=True,
        )
        self.dropout = nn.Dropout(dropout_rate)
        self.pool = nn.AdaptiveMaxPool1d(1)
        self.classifier = nn.Linear(lstm_units * 2, 1)

    def forward(self, input_ids: torch.Tensor,
                attention_mask: torch.Tensor) -> torch.Tensor:
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        h = outputs.last_hidden_state  # (B, T, hidden)
        lstm_out, _ = self.bilstm(h)  # (B, T, 2*lstm_units)
        lstm_out = self.dropout(lstm_out)
        pooled = self.pool(lstm_out.transpose(1, 2)).squeeze(-1)
        return self.classifier(pooled)  # (B, 1)


class BertCNNBiLSTMHead(nn.Module):
    """Encoder -> Conv1d -> ReLU -> BiLSTM -> Dropout -> MaxPool -> Linear.

    Stacked extension of the paper's two heads - Conv1d's local n-gram
    features feed into a BiLSTM for sequence-level integration. Conceptually
    similar to the BERT-CNN-BiLSTM stack used in Mozafari et al. (2020) and
    Faris et al. (2020) for hate speech detection.
    """

    def __init__(
        self,
        model_name: str,
        dropout_rate: float,
        cnn_filters: int = 128,
        cnn_kernel_size: int = 3,
        lstm_units: int = 64,
    ) -> None:
        super().__init__()
        self.bert = _load_encoder(model_name)
        hidden = self.bert.config.hidden_size
        self.cnn = nn.Conv1d(
            in_channels=hidden,
            out_channels=cnn_filters,
            kernel_size=cnn_kernel_size,
            padding=cnn_kernel_size // 2,
        )
        self.bilstm = nn.LSTM(
            input_size=cnn_filters,
            hidden_size=lstm_units,
            batch_first=True,
            bidirectional=True,
        )
        self.dropout = nn.Dropout(dropout_rate)
        self.pool = nn.AdaptiveMaxPool1d(1)
        self.classifier = nn.Linear(lstm_units * 2, 1)

    def forward(self, input_ids: torch.Tensor,
                attention_mask: torch.Tensor) -> torch.Tensor:
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        h = outputs.last_hidden_state  # (B, T, hidden)
        c = self.cnn(h.transpose(1, 2))  # (B, cnn_filters, T)
        c = torch.relu(c)
        lstm_out, _ = self.bilstm(c.transpose(1, 2))  # (B, T, 2*lstm_units)
        lstm_out = self.dropout(lstm_out)
        pooled = self.pool(lstm_out.transpose(1, 2)).squeeze(-1)
        return self.classifier(pooled)  # (B, 1)