"""Shared prefill: encode a state once with the base, then run questions or generation on it.

The state is prefilled with the adapter off, which is exactly what the masked adapter computes
on state tokens. Each question suffix then runs with the adapter on against that cache, and the
cache is cropped back to the state afterwards, so any number of questions, and an escalation to
generation, reuse one prefill.
"""

import numpy as np
import torch

from dmm import lora
from dmm.head import option_logits
from dmm.template import Template, decision_text


class StateCache:
    def __init__(self, model, ids: list[int]):
        self.model, self.ids = model, list(ids)
        with torch.no_grad(), lora.disabled(model):
            out = model.model(input_ids=torch.tensor([ids], device=model.device), use_cache=True)
        self.kv = out.past_key_values

    def __len__(self) -> int:
        return len(self.ids)

    @torch.no_grad()
    def run_suffix(self, suffix: list[int], adapter: bool = True) -> torch.Tensor:
        """Final hidden state at the suffix's last token; the cache is cropped back after."""
        try:
            ctx = lora.disabled(self.model) if not adapter else _null()
            with ctx:
                out = self.model.model(
                    input_ids=torch.tensor([suffix], device=self.model.device),
                    past_key_values=self.kv,
                    use_cache=True,
                )
            return out.last_hidden_state[0, -1]
        finally:
            self.kv.crop(len(self.ids))


class _null:
    def __enter__(self):
        return None

    def __exit__(self, *exc):
        return False


class CachedDecisionModel:
    """decide(items) through the shared-prefill path: one state prefill per distinct state."""

    def __init__(self, model, tokenizer, head, temperatures=None):
        self.model, self.head, self.tpl = model, head, Template(tokenizer)
        self.temperatures = temperatures or {}
        self.prefills = 0

    @torch.no_grad()
    def option_logits(self, items: list[dict]) -> list[torch.Tensor]:
        out: list = [None] * len(items)
        by_state: dict[str, list[int]] = {}
        for i, it in enumerate(items):
            by_state.setdefault(it["state"], []).append(i)
        for _, idx in by_state.items():
            cache = StateCache(self.model, self.tpl.state_ids(items[idx[0]]))
            self.prefills += 1
            for i in idx:
                h = cache.run_suffix(self.tpl.encode(decision_text(items[i])))
                out[i] = option_logits(self.head(h[None]), [items[i]])[0].float().cpu()
        return out

    def decide(self, items: list[dict]) -> list[np.ndarray]:
        out = []
        for z, it in zip(self.option_logits(items), items, strict=True):
            z = z / self.temperatures.get(it["kind"], 1.0)
            out.append(torch.softmax(z, -1).numpy().astype(np.float64))
        return out