# Copyright 2026 The AiVa proof-of-concept team. Licensed under the Apache License, Version 2.0.
"""Load and score an RLCD checkpoint: a ModernBERT encoder, masked mean pooling, one linear head per output,
softmax(logits / temperature). The same class the model card uses; train.py saves in this layout."""

import json
from pathlib import Path

import torch
from safetensors.torch import load_file
from transformers import AutoConfig, AutoModel, AutoTokenizer


def render(messages: list[dict[str, str]]) -> str:
    """The model's input: one `[role] content` line per message, in order."""
    return "\n".join(f"[{m['role']}] {m['content']}" for m in messages)


def pool(hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
    mask = attention_mask.unsqueeze(-1).to(hidden.dtype)
    return (hidden * mask).sum(1) / mask.sum(1).clamp(min=1.0)  # masked mean pooling


class RlcdClassifier(torch.nn.Module):
    def __init__(self, path: Path) -> None:
        super().__init__()
        self.config = json.loads((path / "rlcd_config.json").read_text())
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.encoder = AutoModel.from_config(AutoConfig.from_pretrained(path / "encoder"))
        hidden = self.encoder.config.hidden_size
        self.heads = torch.nn.ModuleDict(
            {name: torch.nn.Linear(hidden, len(classes)) for name, classes in self.config["heads"].items()}
        )
        # Parameter names match the trainer ("encoder.*", "heads.<name>.*"); strict on purpose.
        self.load_state_dict(load_file(str(path / "model.safetensors")), strict=True)
        self.eval()

    render = staticmethod(render)

    def logits(self, texts: list[str]) -> dict[str, torch.Tensor]:
        batch = self.tokenizer(
            texts, padding=True, truncation=True, max_length=self.config["max_length"], return_tensors="pt"
        )
        hidden = self.encoder(**batch).last_hidden_state
        pooled = pool(hidden, batch["attention_mask"])
        return {name: head(pooled) for name, head in self.heads.items()}

    @torch.no_grad()
    def predict(self, messages: list[dict[str, str]]) -> dict[str, dict[str, float]]:
        logits = self.logits([self.render(messages)])
        return {
            name: dict(zip(self.config["heads"][name],
                           (value / self.config["temperatures"][name]).softmax(-1)[0].tolist()))
            for name, value in logits.items()
        }