"""
Cev — Calibrated multi-label intent classifier for HR query routing.

Architecture:
  - Backbone: microsoft/deberta-v3-small (frozen, 141M params)
  - Head: 2-layer MLP with GELU + Dropout → sigmoid output
  - Training loss: Brier score (proper scoring rule — intrinsic calibration)
  - Labels: 4 HR capability buckets (multi-label, non-exclusive)

Usage:
    from transformers import AutoTokenizer
    from modeling_cev import CevModel

    model = CevModel.load("Nehaa/cev-141m")
    tokenizer = AutoTokenizer.from_pretrained("microsoft/deberta-v3-small")
    results = model.predict("What skill gaps does Janet Wise have?", tokenizer)
"""

from __future__ import annotations

from typing import Optional

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer, PretrainedConfig


CAPABILITIES = ["person_lookup", "skill_assessment", "learning_discovery", "assignment"]


class CevConfig(PretrainedConfig):
    model_type = "cev"

    def __init__(
        self,
        backbone: str = "microsoft/deberta-v3-small",
        hidden_size: int = 768,
        mlp_hidden: int = 384,
        dropout: float = 0.1,
        capabilities: list[str] | None = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.backbone = backbone
        self.hidden_size = hidden_size
        self.mlp_hidden = mlp_hidden
        self.dropout = dropout
        self.capabilities = capabilities or CAPABILITIES


class CevModel(nn.Module):
    """
    DeBERTa-v3-small backbone (frozen) + 2-layer MLP head trained with Brier loss.
    Load via CevModel.load("Nehaa/cev-141m").
    """

    def __init__(self, config: CevConfig):
        super().__init__()
        self.config = config
        self.encoder = AutoModel.from_pretrained(config.backbone)
        for param in self.encoder.parameters():
            param.requires_grad = False

        self.head = nn.Sequential(
            nn.Linear(config.hidden_size, config.mlp_hidden),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.mlp_hidden, len(config.capabilities)),
        )

    @classmethod
    def load(cls, repo_id: str = "Nehaa/cev-141m") -> "CevModel":
        """
        Load Cev from Hugging Face Hub.

        Args:
            repo_id: HF repo ID, default "Nehaa/cev-141m".

        Returns:
            CevModel ready for inference.
        """
        from huggingface_hub import hf_hub_download
        import json

        # Load config
        config_path = hf_hub_download(repo_id, "config.json")
        with open(config_path) as f:
            cfg_dict = json.load(f)
        config = CevConfig(**{
            k: v for k, v in cfg_dict.items()
            if k in ("backbone", "hidden_size", "mlp_hidden", "dropout", "capabilities")
        })

        model = cls(config)

        # Load head weights from safetensors
        weights_path = hf_hub_download(repo_id, "model.safetensors")
        from safetensors.torch import load_file
        state = load_file(weights_path)

        # Extract only head weights (ignore encoder keys)
        head_state = {
            k.replace("head.", ""): v
            for k, v in state.items()
            if k.startswith("head.")
        }
        model.head.load_state_dict(head_state)
        model.eval()
        return model

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        token_type_ids: Optional[torch.Tensor] = None,
    ) -> dict[str, torch.Tensor]:
        kwargs = {"input_ids": input_ids, "attention_mask": attention_mask}
        if token_type_ids is not None:
            kwargs["token_type_ids"] = token_type_ids

        with torch.no_grad():
            enc_out = self.encoder(**kwargs)

        cls = enc_out.last_hidden_state[:, 0, :].float()
        logits = self.head(cls)
        probs = torch.sigmoid(logits)
        return {"logits": logits, "probs": probs}

    def predict(
        self,
        text: str | list[str],
        tokenizer: AutoTokenizer,
        threshold: float = 0.5,
        max_length: int = 128,
    ) -> list[dict]:
        """
        Classify one or more queries.

        Args:
            text: String or list of strings.
            tokenizer: AutoTokenizer for microsoft/deberta-v3-small.
            threshold: Decision boundary (default 0.5).
            max_length: Max token length (default 128).

        Returns:
            List of {"probabilities": {cap: float}, "predicted": {cap: 0|1}}
        """
        if isinstance(text, str):
            text = [text]

        device = next(self.parameters()).device
        enc = tokenizer(
            text,
            max_length=max_length,
            padding=True,
            truncation=True,
            return_tensors="pt",
        )
        enc = {k: v.to(device) for k, v in enc.items()}

        self.eval()
        with torch.no_grad():
            out = self.forward(**enc)

        probs = out["probs"].cpu()
        results = []
        for row in probs:
            cap_probs = {c: round(p.item(), 4) for c, p in zip(self.config.capabilities, row)}
            predicted = {c: int(p >= threshold) for c, p in cap_probs.items()}
            results.append({"probabilities": cap_probs, "predicted": predicted})
        return results