"""MiniCPM5 backbone with a scalar option-scoring head for CPM-jev."""

from __future__ import annotations

from pathlib import Path
from typing import Any

import torch
from peft import PeftModel
from safetensors.torch import load_file
from torch import nn
from transformers import AutoModel


BASE_MODEL_ID = "openbmb/MiniCPM5-2B-Base"


class MiniCPMJEVModel(nn.Module):
    """A MiniCPM5 encoder backbone plus a learned scalar candidate scorer."""

    def __init__(self, backbone: nn.Module, hidden_size: int):
        super().__init__()
        self.backbone = backbone
        self.decision_head = nn.Linear(hidden_size, 1, bias=True)

    @staticmethod
    def _hidden_size(config: Any) -> int:
        text_config = config.get_text_config() if hasattr(config, "get_text_config") else config
        for name in ("hidden_size", "dim", "n_embd"):
            if hasattr(text_config, name):
                return int(getattr(text_config, name))
        raise ValueError("Cannot determine the MiniCPM hidden size from its config")

    @classmethod
    def from_pretrained(
        cls,
        model_dir: str | Path,
        *,
        base_model: str = BASE_MODEL_ID,
        dtype: torch.dtype | None = None,
    ) -> "MiniCPMJEVModel":
        model_dir = Path(model_dir)
        if dtype is None:
            dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
        base = AutoModel.from_pretrained(base_model, dtype=dtype, trust_remote_code=True)
        backbone = PeftModel.from_pretrained(base, str(model_dir), is_trainable=False)
        model = cls(backbone, cls._hidden_size(base.config))
        head_path = model_dir / "decision_head.safetensors"
        model.decision_head.load_state_dict(load_file(str(head_path)))
        return model

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        outputs = self.backbone(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=False,
            return_dict=True,
        )
        hidden = outputs.last_hidden_state
        positions = torch.arange(hidden.shape[1], device=hidden.device).unsqueeze(0)
        last = positions.masked_fill(attention_mask.eq(0), -1).max(dim=1).values.clamp_min(0)
        pooled = hidden[torch.arange(hidden.shape[0], device=hidden.device), last]
        return self.decision_head(pooled.to(self.decision_head.weight.dtype)).squeeze(-1)