import torch
import torch.nn as nn

from transformers import AutoModel, PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput


class CustomSequenceClassificationModel(PreTrainedModel):
    """Sequence classification model with a configurable GELU + dropout dense head."""

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config
        self.backbone = AutoModel.from_config(config)

        hidden_size = getattr(config, "hidden_size", None)
        if hidden_size is None:
            raise ValueError(
                "Backbone config is missing 'hidden_size'; cannot build custom head."
            )

        # Primary config keys for Hub reloads.
        # Fallback keeps compatibility with previously saved configs.
        head_layers = getattr(
            config, "head_layers", getattr(config, "classifier_head_layers", 1)
        )
        head_dropout = getattr(
            config, "head_dropout", getattr(config, "classifier_head_dropout", 0.1)
        )
        self.head_layers = int(head_layers)
        self.head_dropout = float(head_dropout)

        self.classifier = self._build_head(
            input_dim=hidden_size,
            output_dim=self.num_labels,
            head_layers=self.head_layers,
            head_dropout=self.head_dropout,
        )
        self.post_init()

    @staticmethod
    def _build_head(
        input_dim: int, output_dim: int, head_layers: int, head_dropout: float
    ) -> nn.Module:
        if head_layers < 1:
            raise ValueError("head_layers must be >= 1")

        # head_layers includes the final projection layer.
        # For head_layers=1 this is equivalent to a single linear classifier.
        if head_layers == 1:
            return nn.Linear(input_dim, output_dim)

        layers = []
        for _ in range(head_layers - 1):
            layers.extend(
                [
                    nn.Linear(input_dim, input_dim),
                    nn.GELU(),
                    nn.Dropout(head_dropout),
                ]
            )
        layers.append(nn.Linear(input_dim, output_dim))
        return nn.Sequential(*layers)

    def _pooled_representation(self, outputs) -> torch.Tensor:
        if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
            return outputs.pooler_output
        return outputs.last_hidden_state[:, 0]

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        labels=None,
        **kwargs,
    ):
        backbone_inputs = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            **kwargs,
        }
        if token_type_ids is not None:
            backbone_inputs["token_type_ids"] = token_type_ids

        backbone_outputs = self.backbone(**backbone_inputs)

        pooled = self._pooled_representation(backbone_outputs)
        logits = self.classifier(pooled)

        loss = None
        if labels is not None:
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif torch.is_floating_point(labels):
                    self.config.problem_type = "multi_label_classification"
                else:
                    self.config.problem_type = "single_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = nn.MSELoss()
                loss = loss_fct(logits.squeeze(), labels.squeeze())
            elif self.config.problem_type == "single_label_classification":
                loss_fct = nn.CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            else:
                loss_fct = nn.BCEWithLogitsLoss()
                loss = loss_fct(logits, labels)

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=backbone_outputs.hidden_states,
            attentions=backbone_outputs.attentions,
        )