"""Jev-style decision model following "Releasing the Inner Jev" (Geuter et al., 2026).

A decision is rendered as a lettered multiple-choice prompt (SemIf system message, JSON user
message, the model's own chat template with thinking disabled). One forward pass gives the final
hidden state ``h`` at the last prompt position, and the decision head maps it to one logit per
option slot:

    z = W0 h + A std(h) + b

``W0`` holds the frozen language-model-head rows for the letters A..P, so the untrained model is the
zero-shot letter readout. ``A`` and ``b`` start at zero and are trained jointly with LoRA adapters.
A softmax over the question's K options (divided by a per-type temperature) gives the probabilities.
"""

from __future__ import annotations

import json
import random
from collections import Counter
from pathlib import Path
from typing import Any

import torch
from torch import nn

SYSTEM_PROMPT = (
    "Apply the supplied criterion to the supplied evidence. Choose exactly one listed option. "
    "Respond with only its uppercase letter, with no explanation or reasoning."
)
LETTERS = "ABCDEFGHIJKLMNOP"
MAX_OPTIONS = len(LETTERS)
QUESTION_TYPES = ("choice", "noul", "score")
TRUNCATION_MARK = " [truncated]"
MASKED_LOGIT = -1e9


def question_type(row: dict[str, Any]) -> str:
    return row["question"]["type"]


def request_order(row: dict[str, Any]) -> list[str]:
    """Evaluation order: choice in request order, yes/no as true then false, score low to high."""
    if question_type(row) == "noul":
        return ["true", "false"]
    return list(row["labels"])


def training_order(row: dict[str, Any], rng: random.Random) -> list[str]:
    """Options are shuffled every presentation; score levels keep their order."""
    order = request_order(row)
    if question_type(row) != "score":
        rng.shuffle(order)
    return order


def user_message(row: dict[str, Any], order: list[str], state: str | None = None) -> str:
    criteria = row["question"]["criteria"]
    criterion = row["question"]["instructions"]
    surface = (row.get("focus") or {}).get("surface")
    if surface:
        criterion = f"{criterion} Focus entity: {surface}."
    return json.dumps({
        "evidence": row["state"] if state is None else state,
        "criterion": criterion,
        "options": [
            {"letter": LETTERS[index], "description": f"{label}: {criteria[label]}"}
            for index, label in enumerate(order)
        ],
    }, ensure_ascii=False)


def _chat_ids(tokenizer: Any, user: str) -> list[int]:
    messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user}]
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
    )
    return tokenizer(text, add_special_tokens=False)["input_ids"]


def encode(
    tokenizer: Any,
    row: dict[str, Any],
    order: list[str],
    max_length: int,
    truncate_state: bool,
) -> tuple[list[int] | None, bool]:
    """Token ids for one decision and whether the state was truncated.

    Over-long inputs are shortened from the end of the state, never the options or the final
    answer position. With ``truncate_state=False`` (training, as in the paper) they return None.
    """
    if len(order) > MAX_OPTIONS:
        raise ValueError(f"{row['id']}: {len(order)} options; the letter readout supports {MAX_OPTIONS}")
    ids = _chat_ids(tokenizer, user_message(row, order))
    if len(ids) <= max_length:
        return ids, False
    if not truncate_state:
        return None, False
    state_ids = tokenizer(row["state"], add_special_tokens=False)["input_ids"]
    overflow = len(ids) - max_length
    for _ in range(4):
        keep = len(state_ids) - overflow - 8
        if keep <= 0:
            break
        state = tokenizer.decode(state_ids[:keep]) + TRUNCATION_MARK
        ids = _chat_ids(tokenizer, user_message(row, order, state))
        if len(ids) <= max_length:
            return ids, True
        overflow += len(ids) - max_length
    raise ValueError(f"{row['id']}: question and options alone exceed max_length={max_length}")


def letter_token_ids(tokenizer: Any) -> list[int]:
    ids = []
    for letter in LETTERS:
        token_ids = tokenizer.encode(letter, add_special_tokens=False)
        if len(token_ids) != 1:
            raise ValueError(f"letter {letter!r} is not a single token for this tokenizer: {token_ids}")
        ids.append(token_ids[0])
    return ids


class DecisionHead(nn.Module):
    """z = W0 h + A std(h) + b with W0 frozen from the LM head and A, b initialised to zero."""

    def __init__(self, letter_rows: torch.Tensor, letter_bias: torch.Tensor | None = None,
                 softcap: float | None = None):
        super().__init__()
        slots, hidden = letter_rows.shape
        self.register_buffer("letter_rows", letter_rows.detach().float().clone())
        self.register_buffer(
            "letter_bias",
            torch.zeros(slots) if letter_bias is None else letter_bias.detach().float().clone(),
        )
        self.register_buffer("feature_mean", torch.zeros(hidden))
        self.register_buffer("feature_std", torch.ones(hidden))
        self.correction = nn.Parameter(torch.zeros(slots, hidden))
        self.bias = nn.Parameter(torch.zeros(slots))
        self.softcap = softcap

    def set_feature_stats(self, features: torch.Tensor) -> None:
        """Fix the standardisation once on training features, before training."""
        features = features.float()
        self.feature_mean.copy_(features.mean(dim=0))
        self.feature_std.copy_(features.std(dim=0, unbiased=False).clamp(min=1e-6))

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        hidden = hidden.float()
        readout = hidden @ self.letter_rows.T + self.letter_bias
        if self.softcap:
            readout = self.softcap * torch.tanh(readout / self.softcap)
        standardized = (hidden - self.feature_mean) / self.feature_std
        return readout + standardized @ self.correction.T + self.bias


class DecisionModel(nn.Module):
    """Decoder backbone (optionally LoRA-wrapped) plus the letter-initialised decision head."""

    def __init__(self, backbone: nn.Module, head: DecisionHead):
        super().__init__()
        self.backbone = backbone
        self.head = head

    def last_hidden(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        hidden = self.backbone(
            input_ids=input_ids, attention_mask=attention_mask, use_cache=False
        ).last_hidden_state
        last = attention_mask.sum(dim=1).clamp(min=1) - 1
        pooled = hidden[torch.arange(hidden.size(0), device=hidden.device), last.to(hidden.device)]
        return pooled.to(self.head.letter_rows.device)

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor,
                option_counts: torch.Tensor) -> torch.Tensor:
        logits = self.head(self.last_hidden(input_ids, attention_mask))
        slots = torch.arange(logits.size(1), device=logits.device)
        invalid = slots.unsqueeze(0) >= option_counts.to(logits.device).unsqueeze(1)
        return logits.masked_fill(invalid, MASKED_LOGIT)


def pad_batch(sequences: list[list[int]], pad_id: int, device: torch.device | str):
    width = max(len(sequence) for sequence in sequences)
    input_ids = torch.full((len(sequences), width), pad_id, dtype=torch.long)
    attention_mask = torch.zeros((len(sequences), width), dtype=torch.long)
    for index, sequence in enumerate(sequences):
        input_ids[index, :len(sequence)] = torch.tensor(sequence, dtype=torch.long)
        attention_mask[index, :len(sequence)] = 1
    return input_ids.to(device), attention_mask.to(device)


def parse_max_memory(spec: str | None) -> dict[Any, str] | None:
    if not spec:
        return None
    result: dict[Any, str] = {}
    for part in spec.split(","):
        key, value = part.split("=")
        result[int(key) if key.strip().isdigit() else key.strip()] = value.strip()
    return result


def pick_dtype(counter: Counter) -> torch.dtype:
    if not torch.cuda.is_available():
        counter["cpu_execution_float32"] += 1
        return torch.float32
    if all(torch.cuda.get_device_capability(index)[0] >= 8 for index in range(torch.cuda.device_count())):
        return torch.bfloat16
    counter["float16_no_bf16_device"] += 1
    return torch.float16


def load_decision_model(
    model_id: str,
    *,
    quantization: str = "none",
    max_memory: str | None = None,
    adapter_dir: Path | None = None,
    head_path: Path | None = None,
    trainable_adapter: bool = False,
    counter: Counter | None = None,
    local_files_only: bool = False,
):
    """Load backbone, tokenizer and head. Returns (model, tokenizer, letter_ids, dtype)."""
    from transformers import AutoModelForCausalLM, AutoTokenizer

    counter = counter if counter is not None else Counter()
    tokenizer = AutoTokenizer.from_pretrained(model_id, local_files_only=local_files_only)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
        counter["pad_token_from_eos"] += 1
    dtype = pick_dtype(counter)
    if torch.cuda.is_available() and any(
        torch.cuda.get_device_capability(index)[0] < 8 for index in range(torch.cuda.device_count())
    ):
        # The flash SDPA kernel needs Ampere+ (the RTX 2070 SUPER is Turing); keep mem-efficient SDPA.
        torch.backends.cuda.enable_flash_sdp(False)
        counter["flash_sdp_disabled_pre_ampere_gpu"] += 1
    kwargs: dict[str, Any] = {"dtype": dtype, "local_files_only": local_files_only,
                              "attn_implementation": "sdpa"}
    if torch.cuda.is_available():
        kwargs["device_map"] = "auto"
        memory = parse_max_memory(max_memory)
        if memory:
            kwargs["max_memory"] = memory
    if quantization == "4bit":
        from transformers import BitsAndBytesConfig
        kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=dtype,
        )
    elif quantization != "none":
        raise ValueError(f"unknown quantization {quantization!r}")

    language_model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
    letter_ids = letter_token_ids(tokenizer)
    offloaded = sorted({str(d) for d in (getattr(language_model, "hf_device_map", None) or {}).values()
                        if str(d) in ("cpu", "disk")})
    if offloaded:
        counter["layers_offloaded_to_" + "_".join(offloaded)] += 1
    output = language_model.get_output_embeddings()
    weight = output.weight
    if weight.device.type == "meta":
        # Offloaded by accelerate: the real tensor lives in the module hook's weights map.
        weight = output._hf_hook.weights_map["weight"]
    letter_rows = weight[letter_ids].detach().float().cpu()
    letter_bias = output.bias[letter_ids].detach().float().cpu() if getattr(output, "bias", None) is not None else None
    config = language_model.config
    softcap = getattr(config, "final_logit_softcapping", None) or getattr(
        getattr(config, "text_config", None), "final_logit_softcapping", None
    )
    backbone = language_model.get_decoder()
    del language_model

    head = DecisionHead(letter_rows, letter_bias, softcap)
    if head_path is not None:
        payload = torch.load(head_path, map_location="cpu", weights_only=True)
        if payload["letter_ids"] != letter_ids:
            raise ValueError("checkpoint letter ids do not match this tokenizer")
        head.load_state_dict(payload["head_state_dict"])
    head_device = next(backbone.parameters()).device
    head.to(head_device)

    if adapter_dir is not None:
        from peft import PeftModel
        backbone = PeftModel.from_pretrained(backbone, str(adapter_dir), is_trainable=trainable_adapter)
    model = DecisionModel(backbone, head)
    return model, tokenizer, letter_ids, dtype


def option_logits(model: DecisionModel, tokenizer: Any, row: dict[str, Any], max_length: int,
                  counter: Counter) -> tuple[dict[str, float], bool]:
    """Raw decision-head logits for the row's options, in request order (one forward pass)."""
    order = request_order(row)
    ids, truncated = encode(tokenizer, row, order, max_length, truncate_state=True)
    if truncated:
        counter["state_truncated"] += 1
    device = next(model.backbone.parameters()).device
    input_ids, attention_mask = pad_batch([ids], tokenizer.pad_token_id, device)
    counts = torch.tensor([len(order)])
    model.eval()
    with torch.inference_mode():
        logits = model(input_ids, attention_mask, counts)[0, :len(order)].float().cpu()
    return {label: float(logits[index]) for index, label in enumerate(order)}, truncated


def probabilities(logits: dict[str, float], temperature: float = 1.0) -> dict[str, float]:
    values = torch.tensor(list(logits.values()), dtype=torch.float64) / temperature
    probs = torch.softmax(values, dim=0).tolist()
    return dict(zip(logits, probs))


def jev_answer(row: dict[str, Any], probs: dict[str, float]) -> dict[str, Any]:
    """Typed answer in TypeSafe's format: choice, noul (P(true)) or score."""
    kind = question_type(row)
    best = max(probs, key=probs.get)
    if kind == "noul":
        return {"type": "noul", "noul": probs["true"]}
    if kind == "score":
        return {"type": "score", "score": int(best) if best.isdigit() else best, "probabilities": probs}
    return {"type": "choice", "choice": best, "probabilities": probs}