import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModel, PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
from .configuration_baseline_enc import BaselineEncConfig

class Encoder(nn.Module):
    def __init__(self, name, pooling="cls"):
        super().__init__()
        base_cfg = AutoConfig.from_pretrained(name)
        self.backbone = AutoModel.from_config(base_cfg)
        self.pooling = pooling
        self.hidden_dim = self.backbone.config.hidden_size

    def forward(self, ids, mask):
        hidden = self.backbone(input_ids=ids, attention_mask=mask).last_hidden_state
        if self.pooling == "cls":
            return hidden[:, 0, :]
        m = mask.unsqueeze(-1).float()
        return (hidden * m).sum(1) / m.sum(1).clamp(min=1e-9)

class SarcasmHead(nn.Module):
    def __init__(self, in_dim, hidden_dim, dropout=0.1):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, hidden):
        return self.mlp(hidden).squeeze(-1)

class BaselineEncForSequenceClassification(PreTrainedModel):
    config_class = BaselineEncConfig
    base_model_prefix = "baseline_enc"

    def __init__(self, config):
        super().__init__(config)
        self.encoder = Encoder(config.encoder_name, config.pooling)
        self.head = SarcasmHead(
            self.encoder.hidden_dim,
            config.mlp_hidden,
            config.dropout,
        )
        self.post_init()

    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
        pooled = self.encoder(input_ids, attention_mask)
        binary_logit = self.head(pooled)
        logits = torch.stack([-binary_logit, binary_logit], dim=-1)

        loss = None
        if labels is not None:
            loss = nn.CrossEntropyLoss()(logits, labels)

        return SequenceClassifierOutput(loss=loss, logits=logits)