# ─────────────────────────────────────────────────────────────────────────────
# Fine-Tuned BERT Regressor Architecture
# src/models/bert_regressor.py
# ─────────────────────────────────────────────────────────────────────────────
import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer


class BertRegressorModel(nn.Module):
    """
    End-to-end fine-tuned BERT Regression model for Extraversion prediction.

    Architecture:
        bert-base-uncased  →  [CLS] hidden state (768-dim)
                           →  Dropout(0.3)
                           →  Linear(768, 256)
                           →  ReLU
                           →  Dropout(0.2)
                           →  Linear(256, 1)
                           →  Sigmoid * 99   (output in 0–99 range)
    """

    def __init__(self, dropout1: float = 0.3, dropout2: float = 0.2):
        super().__init__()
        self.bert = BertModel.from_pretrained("bert-base-uncased")
        hidden = self.bert.config.hidden_size   # 768

        self.regressor = nn.Sequential(
            nn.Dropout(dropout1),
            nn.Linear(hidden, 256),
            nn.ReLU(),
            nn.Dropout(dropout2),
            nn.Linear(256, 1),
            nn.Sigmoid()          # → (0, 1); we multiply by 99 at inference
        )

    def forward(self, input_ids, attention_mask, token_type_ids=None):
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            output_attentions=True      # expose attention weights for heatmap
        )
        cls_vec     = outputs.last_hidden_state[:, 0, :]   # [B, 768]
        attentions  = outputs.attentions                   # tuple of 12 layers
        score       = self.regressor(cls_vec).squeeze(-1)  # [B]
        score       = score * 99.0                         # scale to 0–99
        return score, attentions


def get_tokenizer():
    return BertTokenizer.from_pretrained("bert-base-uncased")


def load_bert_regressor(model_path: str, device: str = "cpu") -> BertRegressorModel:
    """Load a saved BertRegressorModel from disk."""
    model = BertRegressorModel()
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()
    model.to(device)
    return model