"""RubriX: score a response against a rubric's own level descriptions, without text generation.

    from transformers.dynamic_module_utils import get_class_from_dynamic_module
    RubriX = get_class_from_dynamic_module("rubrix.RubriX", "CobuildX-AI/RubriX-8B")
    judge = RubriX.from_pretrained("CobuildX-AI/RubriX-8B")          # needs a GPU with ~20 GB (bf16)
    result = judge.score(
        criterion="Does the response answer the question correctly and completely?",
        levels=["Wrong or unrelated.", "Partly correct.", "Correct but incomplete.", "Correct and complete."],
        response="Paris is the capital of France.",
        conversation="What is the capital of France?",
    )
    result["probabilities"]     # one probability per level, lowest level first
    result["level"]             # index of the most likely level (0-based)
    result["expected_level"]    # probability-weighted level (0-based), useful for ranking

How it works: one encoder (Qwen3-8B with this repository's LoRA adapter) embeds the input and each level description
separately; two projection heads map them into a shared space; the score of level k is τ · cos(f(input), g(level_k))
and a softmax over the rubric's levels gives the probabilities. Level vectors do not depend on the input, so they are
cached: judging many responses on one rubric embeds its levels once.

The input text is laid out exactly as in training:

    [system prompt]          (optional)
    ...
    [conversation]           (optional: the prompt or dialogue the response answers)
    ...
    [reference answer]       (optional)
    ...
    [response to evaluate]
    ...

    <criterion>

Requirements: torch, transformers, peft, safetensors, huggingface_hub.
"""
from __future__ import annotations

import json
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F

BASE_MODEL = "Qwen/Qwen3-8B"
BASE_REVISION = "b968826d9c46dd6066d109eabc6255188de91218"   # the exact revision RubriX was trained on


class _Head(nn.Module):
    """hidden -> width -> width -> proj MLP with GELU and LayerNorm (the CLM head layout)."""

    def __init__(self, hidden: int, width: int, depth: int, proj: int, layernorm: bool):
        super().__init__()
        self.inp = nn.Linear(hidden, width)
        self.hidden = nn.ModuleList(nn.Linear(width, width) for _ in range(depth - 2))
        self.norms = nn.ModuleList((nn.LayerNorm(width) if layernorm else nn.Identity()) for _ in range(depth - 2))
        self.out = nn.Linear(width, proj)
        self.act = nn.GELU()

    def forward(self, x):
        x = self.act(self.inp(x))
        for lin, nrm in zip(self.hidden, self.norms):
            x = self.act(nrm(lin(x)))
        return self.out(x)


def build_input(response: str, criterion: str, conversation: str | None = None, system_prompt: str | None = None,
                reference_answer: str | None = None) -> str:
    """The text the input encoder sees: labelled sections, then the criterion after a blank line."""
    parts = {"system prompt": system_prompt, "conversation": conversation, "reference answer": reference_answer,
             "response to evaluate": response}
    state = "\n\n".join(f"[{k}]\n{v}" for k, v in parts.items() if isinstance(v, str) and v.strip())
    s, c = state.strip(), (criterion or "").strip()
    return f"{s}\n\n{c}" if s and c else (s or c)


class RubriX:
    def __init__(self, tokenizer, encoder, state_head, level_head, scale: float, device, max_tokens: int = 32768):
        self.tok, self.encoder = tokenizer, encoder
        self.state_head, self.level_head = state_head, level_head
        self.scale, self.device, self.max_tokens = scale, device, max_tokens
        self._level_cache: dict[str, torch.Tensor] = {}

    @classmethod
    def from_pretrained(cls, repo_id: str = "CobuildX-AI/RubriX-8B", device: str | None = None,
                        dtype: torch.dtype = torch.bfloat16, merge: bool = True, token: str | bool | None = None,
                        max_tokens: int | None = None):
        """Load Qwen3-8B (pinned revision), apply the LoRA adapter and load the heads. `repo_id` may be a local folder.
        `max_tokens` caps each input (default 32,768, Qwen3-8B's native context; RubriX was trained on up to 8,192)."""
        from huggingface_hub import snapshot_download
        from peft import PeftModel
        from safetensors.torch import load_file
        from transformers import AutoModel, AutoTokenizer

        path = Path(repo_id) if Path(repo_id).is_dir() else Path(snapshot_download(repo_id, token=token))
        cfg = json.loads((path / "rubrix_config.json").read_text())
        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        tok = AutoTokenizer.from_pretrained(cfg["base_model"], revision=cfg["base_revision"])
        tok.padding_side = "right"                    # pooling reads the last real token of each row
        base = AutoModel.from_pretrained(cfg["base_model"], revision=cfg["base_revision"], dtype=dtype)
        enc = PeftModel.from_pretrained(base, str(path / "adapter"))
        enc = enc.merge_and_unload() if merge else enc
        enc = enc.to(device).eval()
        h = cfg["heads"]
        heads = load_file(str(path / "heads.safetensors"))
        state_head = _Head(h["hidden_size"], h["width"], h["depth"], h["projection_dim"], h["layernorm"])
        level_head = _Head(h["hidden_size"], h["width"], h["depth"], h["projection_dim"], h["layernorm"])
        state_head.load_state_dict({k[len("state_head."):]: v for k, v in heads.items() if k.startswith("state_head.")})
        level_head.load_state_dict({k[len("level_head."):]: v for k, v in heads.items() if k.startswith("level_head.")})
        scale = min(float(torch.exp(heads["logit_scale"].float())), h.get("max_scale", 100.0))
        return cls(tok, enc, state_head.to(device).float().eval(), level_head.to(device).float().eval(), scale, device,
                   max_tokens or cfg.get("max_tokens", 32768))

    @torch.no_grad()
    def embed(self, texts: list[str], batch_size: int = 8) -> torch.Tensor:
        """L2-normalised final hidden state of each text's last token."""
        out = []
        for i in range(0, len(texts), batch_size):
            b = self.tok(texts[i:i + batch_size], return_tensors="pt", padding=True, truncation=True,
                         max_length=self.max_tokens).to(self.device)
            h = self.encoder(**b, use_cache=False).last_hidden_state
            last = b["attention_mask"].sum(1) - 1
            out.append(F.normalize(h[torch.arange(h.shape[0], device=h.device), last].float(), dim=-1))
        return torch.cat(out)

    @torch.no_grad()
    def _levels(self, levels: list[str]) -> torch.Tensor:
        new = [t for t in dict.fromkeys(levels) if t not in self._level_cache]
        if new:
            for t, v in zip(new, F.normalize(self.level_head(self.embed(new)), dim=-1)):
                self._level_cache[t] = v
        return torch.stack([self._level_cache[t] for t in levels])

    def clear_cache(self):
        self._level_cache.clear()

    @torch.no_grad()
    def score_batch(self, items: list[dict], batch_size: int = 8) -> list[dict]:
        """items: dicts with criterion, levels (lowest first), response, and optional conversation, system_prompt,
        reference_answer. Returns one result per item."""
        texts = [build_input(it["response"], it["criterion"], it.get("conversation"), it.get("system_prompt"),
                             it.get("reference_answer")) for it in items]
        z = F.normalize(self.state_head(self.embed(texts, batch_size)), dim=-1)
        results = []
        for zi, it in zip(z, items):
            if len(it["levels"]) < 2:
                raise ValueError("a rubric needs at least two levels")
            logits = self.scale * (self._levels(list(it["levels"])) @ zi)
            p = torch.softmax(logits, dim=-1)
            results.append({"probabilities": p.tolist(), "level": int(p.argmax()),
                            "expected_level": float((p * torch.arange(len(p), device=p.device)).sum())})
        return results

    def score(self, criterion: str, levels: list[str], response: str, conversation: str | None = None,
              system_prompt: str | None = None, reference_answer: str | None = None) -> dict:
        """Score one response. `levels` are the rubric's level descriptions, lowest first."""
        return self.score_batch([{"criterion": criterion, "levels": levels, "response": response,
                                  "conversation": conversation, "system_prompt": system_prompt,
                                  "reference_answer": reference_answer}])[0]