"""Amber 0.4B as a Decision Index engine (apolinario/decision-index, edition 0.2.1).

ModernBERT-large with a LoRA adapter (r=8 on Wqkv / Wo / Wi) and an option scorer. The
context (state + question) and each option are encoded apart and mean-pooled; the scorer,
an MLP over [h_ctx; h_opt; h_ctx * h_opt], gives one logit per option and the answer is the
softmax over them. A ``noul`` question is a choice between its ``true`` and ``false``
criteria (``yes`` / ``no`` when it has none), answered with P(true). No abstention.

    python -m decision_index run --engine amber_engine:AmberEngine \
        --option weights=amber.safetensors --suite-dir <suite> --out <run dir> --compact

One rendering for every benchmark: a string state as is, an object state as ``key: value``
lines (non-string values as JSON), each option as its description (its key when there is
none). A context over ``max_tokens`` (ModernBERT's 8,192 positions) is read in consecutive
windows of its longest field (state or question), each with the other field; its vector is
the token-weighted mean of the window vectors. Nothing is truncated; an option over the
limit is unsupported.
"""

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any

import torch
from decision_index.engines import Engine, Unsupported
from peft import LoraConfig, get_peft_model
from safetensors.torch import load_file
from torch import nn
from transformers import AutoModel, AutoTokenizer

BASE = "answerdotai/ModernBERT-large"
BASE_REVISION = "45bb4654a4d5aaff24dd11d4781fa46d39bf8c13"
LORA = {
    "r": 8,
    "lora_alpha": 16,
    "lora_dropout": 0.05,
    "bias": "none",
    "target_modules": ["Wqkv", "Wo", "Wi"],
}
SCORER_WIDTH = 512
DTYPES = {"fp32": torch.float32, "bf16": torch.bfloat16}
NOUL_KEYS = ["true", "false"]
NOUL_OPTIONS = ["yes", "no"]


def render(value: Any) -> str:
    if value is None or value in ("", {}, []):
        return ""
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        return "\n".join(
            f"{k}: {v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)}"
            for k, v in value.items()
        )
    return json.dumps(value, ensure_ascii=False)


def context_text(fields: dict[str, str]) -> str:
    return f"STATE:\n{fields['state']}\n\nQUESTION:\n{fields['question']}"


def option_text(option: str) -> str:
    return f"OPTION:\n{option}"


def masked_mean(h: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    m = mask.unsqueeze(-1).to(h.dtype)
    return (h * m).sum(dim=1) / m.sum(dim=1).clamp(min=1.0)


class OptionScorer(nn.Module):
    """score = MLP([h_ctx; h_opt; h_ctx * h_opt])."""

    def __init__(self, hidden: int, width: int) -> None:
        super().__init__()
        self.net = nn.Sequential(nn.Linear(3 * hidden, width), nn.GELU(), nn.Linear(width, 1))

    def forward(self, h_ctx: torch.Tensor, h_opt: torch.Tensor) -> torch.Tensor:
        if h_ctx.dim() == 1:
            h_ctx = h_ctx.unsqueeze(0)
        if h_ctx.size(0) == 1 and h_opt.size(0) != 1:
            h_ctx = h_ctx.expand(h_opt.size(0), -1)
        return self.net(torch.cat([h_ctx, h_opt, h_ctx * h_opt], dim=-1)).squeeze(-1)


def questions_of(
    state: Any, questions: dict[str, dict]
) -> list[tuple[str, str, dict, list[str], list[str]]]:
    """(question key, type, context fields, options, answer key per option) per question."""
    context = render(state)
    out = []
    for qkey, q in questions.items():
        if q["type"] == "choice":
            keys = list(q["criteria"])
            options = [render(desc) or key for key, desc in q["criteria"].items()]
        elif q["type"] == "noul":
            criteria = q.get("criteria") or {}
            keys = list(NOUL_KEYS)
            options = [render(criteria.get(k)) or o for k, o in zip(NOUL_KEYS, NOUL_OPTIONS)]
        else:
            raise Unsupported(f"unsupported question type {q['type']!r}")
        if not options:
            raise Unsupported(f"question {qkey!r} has no options")
        fields = {"state": context, "question": render(q.get("instructions", ""))}
        out.append((qkey, q["type"], fields, options, keys))
    return out


def answer(q_type: str, keys: list[str], probs: list[float]) -> dict:
    if q_type == "noul":
        return {"type": "noul", "noul": float(probs[0])}
    best = max(range(len(keys)), key=probs.__getitem__)
    return {
        "type": "choice",
        "choice": keys[best],
        "probabilities": dict(zip(keys, map(float, probs))),
    }


class AmberEngine(Engine):
    name = "amber"
    latency = (
        "Device-synchronized in-process request wall time: tokenization, one encoder pass per "
        "question, option encoding for strings not yet cached; excludes model loading."
    )

    def __init__(
        self,
        weights: str = "amber.safetensors",
        device: str | None = None,
        dtype: str = "bf16",
        scorer_dtype: str = "fp32",
        max_tokens: int = 8192,
        cache_size: int = 50_000,
        # Padded tokens per encoder batch; peak memory grows with it.
        batch_tokens: int = 8192,
        **options,
    ):
        super().__init__(**options)
        if dtype not in DTYPES or scorer_dtype not in DTYPES:
            raise ValueError(f"dtype and scorer_dtype must be one of {sorted(DTYPES)}")
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        path = Path(weights)
        state = load_file(str(path))
        self.tok = AutoTokenizer.from_pretrained(BASE, revision=BASE_REVISION)
        backbone = AutoModel.from_pretrained(BASE, revision=BASE_REVISION)
        self.encoder = get_peft_model(backbone, LoraConfig(**LORA))
        lora = {k[len("encoder.") :]: v for k, v in state.items() if k.startswith("encoder.")}
        want = {k for k in self.encoder.state_dict() if ".lora_" in k}
        _, unexpected = self.encoder.load_state_dict(lora, strict=False)
        if unexpected or want - set(lora):
            raise ValueError(
                f"{path}: LoRA keys do not match the adapter ({len(want - set(lora))} missing)"
            )
        self.scorer = OptionScorer(backbone.config.hidden_size, SCORER_WIDTH)
        self.scorer.load_state_dict(
            {k[len("scorer.") :]: v for k, v in state.items() if k.startswith("scorer.")}
        )
        # Option logits sit far from 0 (|logit| ~ 100): bf16 spaces them 0.5-1 apart, so close
        # options would tie and the first one win.
        self.encoder.to(self.device, DTYPES[dtype]).eval()
        self.scorer.to(self.device, DTYPES[scorer_dtype]).eval()
        self.scorer_dtype = DTYPES[scorer_dtype]
        self.dtype = dtype
        self.limit = int(max_tokens)
        self.cache_size = int(cache_size)
        self.batch_tokens = int(batch_tokens)
        self.cache: dict[str, torch.Tensor] = {}
        self.provenance = {
            "kind": "amber",
            "weights": str(path),
            "weights_sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
            "base": f"{BASE}@{BASE_REVISION}",
            "dtype": dtype,
            "scorer_dtype": scorer_dtype,
            "context_limit_tokens": self.limit,
            "policy": (
                "One fixed rendering for every benchmark: string state as is, object state as key: value "
                "lines, option = its description or key. Softmax over the option scorer's logits; noul = "
                "P(true) over its true / false criteria. No abstention, no truncation, no option filtering. "
                "A context over the limit is read in consecutive windows of its longest field (state or "
                "question), each with the other field, and scored on the token-weighted mean of the window "
                "vectors; an over-length option is unsupported."
            ),
        }

    def runtime(self):
        info = {"torch": torch.__version__, "device": self.device, "dtype": self.dtype}
        if self.device == "cuda":
            info.update(cuda=torch.version.cuda, gpu=torch.cuda.get_device_name())
        return info

    def synchronize(self):
        if self.device == "cuda":
            torch.cuda.synchronize()

    def _ids(self, text: str) -> list[int]:
        ids = self.tok(text)["input_ids"]
        if len(ids) > self.limit:
            raise Unsupported(f"input of {len(ids)} tokens over the {self.limit}-token limit")
        return ids

    def _windows(self, fields: dict[str, str]) -> list[list[int]]:
        ids = self.tok(context_text(fields))["input_ids"]
        if len(ids) <= self.limit:
            return [ids]
        field = max(("state", "question"), key=lambda f: len(fields[f]))
        text = fields[field]
        spans = self.tok(text, add_special_tokens=False, return_offsets_mapping=True)[
            "offset_mapping"
        ]
        # Margin for tokens that merge differently at a window's edges.
        budget = self.limit - (len(ids) - len(spans)) - 16
        if budget < 512:
            raise Unsupported(f"input of {len(ids)} tokens: the other fields fill the limit")
        out = []
        for s in range(0, len(spans), budget):
            piece = text[spans[s][0] : spans[min(s + budget, len(spans)) - 1][1]]
            out.append(self._ids(context_text({**fields, field: piece})))
        return out

    def _encode(self, batch: list[list[int]]) -> torch.Tensor:
        ids = nn.utils.rnn.pad_sequence(
            [torch.tensor(x) for x in batch], batch_first=True, padding_value=0
        )
        mask = nn.utils.rnn.pad_sequence(
            [torch.ones(len(x), dtype=torch.long) for x in batch], batch_first=True
        )
        ids, mask = ids.to(self.device), mask.to(self.device)
        return masked_mean(self.encoder(input_ids=ids, attention_mask=mask).last_hidden_state, mask)

    def _encode_all(self, ids: list[list[int]]) -> list[torch.Tensor]:
        out: list[torch.Tensor] = [None] * len(ids)  # type: ignore[list-item]
        batch: list[int] = []
        # Shortest first, so a batch's padded size is (rows) x (its last row's length).
        for i in sorted(range(len(ids)), key=lambda i: len(ids[i])):
            if batch and (len(batch) + 1) * len(ids[i]) > self.batch_tokens:
                for j, h in zip(batch, self._encode([ids[j] for j in batch])):
                    out[j] = h
                batch = []
            batch.append(i)
        if batch:
            for j, h in zip(batch, self._encode([ids[j] for j in batch])):
                out[j] = h
        return out

    def _options(self, options: list[str]) -> torch.Tensor:
        new = [o for o in dict.fromkeys(options) if o not in self.cache]
        self.cache.update(zip(new, self._encode_all([self._ids(option_text(o)) for o in new])))
        out = torch.stack([self.cache[o] for o in options])
        while len(self.cache) > self.cache_size:
            self.cache.pop(next(iter(self.cache)))
        return out

    @torch.inference_mode()
    def __call__(self, state, questions):
        rows = questions_of(state, questions)
        windows = [self._windows(fields) for _, _, fields, _, _ in rows]
        encoded = iter(self._encode_all([ids for w in windows for ids in w]))
        answers = {}
        for (qkey, q_type, _, options, keys), w in zip(rows, windows):
            h = torch.stack([next(encoded) for _ in w])
            if len(w) > 1:
                n = torch.tensor([len(ids) for ids in w], device=h.device, dtype=torch.float32)
                h = ((n[:, None] * h.float()).sum(0) / n.sum()).to(h.dtype)[None]
            logits = self.scorer(
                h[0].to(self.scorer_dtype), self._options(options).to(self.scorer_dtype)
            )
            answers[qkey] = answer(q_type, keys, torch.softmax(logits.float(), dim=0).tolist())
        return {"model": "amber", "answers": answers}, None