"""A standalone port of DecidaBERT-large's inference (the decision heads and the Laya-format prompt layout).

This intentionally does not import the `decida` package: it is independent code for a Gradio Space, reading only
the public checkpoint at https://huggingface.co/helmo/DecidaBERT-large. The algorithm matches
`src/decida/model/enc.py` in the main project, ported here because ZeroGPU Spaces only support the Gradio SDK, not
a plain server process like `decida serve`. If the real project's encoder logic changes, this file needs updating
by hand; it is not generated from it.
"""
from __future__ import annotations

import json
import os
from dataclasses import dataclass
from enum import Enum
from typing import Any

import numpy as np
import torch
from safetensors.torch import load_file
from torch import nn
from transformers import AutoConfig, AutoModel, AutoTokenizer

MODEL_REPO = "helmo/DecidaBERT-large"


class Primitive(str, Enum):
    CHOICE = "choice"
    SCORE = "score"
    NOUL = "noul"


@dataclass
class Question:
    qid: str
    primitive: Primitive
    instructions: str
    criteria: Any


@dataclass
class Encoded:
    input_ids: torch.Tensor
    attention_mask: torch.Tensor
    readout_positions: torch.Tensor
    n_options: int
    truncated: bool


@dataclass
class Batch:
    input_ids: torch.Tensor
    attention_mask: torch.Tensor
    readout_positions: torch.Tensor
    marker_mask: torch.Tensor
    qtypes: torch.Tensor


def _render_criterion(value: Any) -> str:
    if isinstance(value, str):
        return value
    return json.dumps(value, ensure_ascii=False, separators=(", ", ": "), default=str)


def _render_options(q: Question) -> list[str]:
    if q.primitive == Primitive.CHOICE:
        return [k if not v else f"{k}: {_render_criterion(v)}" for k, v in q.criteria.items()]
    if q.primitive == Primitive.SCORE:
        return [f"level {i}: {_render_criterion(c)}" for i, c in enumerate(q.criteria)]
    crit = q.criteria or {}
    false_str = _render_criterion(crit.get("false")) if crit.get("false") else "no, the statement does not hold"
    true_str = _render_criterion(crit.get("true")) if crit.get("true") else "yes, the statement holds"
    return [f"false: {false_str}", f"true: {true_str}"]


class LayaLayout:
    """Builds one [CLS] instructions [SEP] <mask>opt1 <mask>opt2 ... [SEP] state [SEP] sequence per question."""

    def __init__(self, tokenizer, max_len: int = 2142, head_max_len: int = 1400, opt_max_tokens: int = 160):
        self.tok, self.max_len, self.head_max_len, self.opt_max_tokens = tokenizer, max_len, head_max_len, opt_max_tokens

    def build(self, state: str, q: Question) -> Encoded:
        mask_tok = self.tok.mask_token
        opts = _render_options(q)
        ins = q.instructions if isinstance(q.instructions, str) else json.dumps(q.instructions, ensure_ascii=False)
        if mask_tok:
            ins = ins.replace(mask_tok, " ")
        head_ids = self.tok(f"{q.primitive.value} question: {ins}", add_special_tokens=False)["input_ids"]

        opt_ids = []
        for o in opts:
            text_ids = self.tok(" " + o.replace(mask_tok, " "), add_special_tokens=False)["input_ids"]
            opt_ids.append([self.tok.mask_token_id] + text_ids[: self.opt_max_tokens])

        opt_budget = self.head_max_len - sum(len(o) for o in opt_ids)
        if opt_budget < 16:
            per = max(4, (self.head_max_len - 16) // max(1, len(opt_ids)))
            opt_ids = [o[:per] for o in opt_ids]
            opt_budget = self.head_max_len - sum(len(o) for o in opt_ids)

        head_ids = head_ids[: max(8, opt_budget)]
        ids = [self.tok.cls_token_id] + head_ids + [self.tok.sep_token_id]
        markers = []
        for o in opt_ids:
            markers.append(len(ids))
            ids.extend(o)
        ids.append(self.tok.sep_token_id)

        room = max(0, self.max_len - len(ids) - 1)
        state_str = json.dumps(state, ensure_ascii=False) if isinstance(state, (dict, list)) else str(state)
        st = self.tok(state_str.replace(mask_tok, " "), add_special_tokens=False)["input_ids"]
        truncated = len(st) > room
        ids = ids + st[:room] + [self.tok.sep_token_id]

        final_ids = ids[: self.max_len]
        valid_markers = [m for m in markers if m < self.max_len]
        return Encoded(
            input_ids=torch.tensor(final_ids, dtype=torch.long),
            attention_mask=torch.ones(len(final_ids), dtype=torch.long),
            readout_positions=torch.tensor(valid_markers, dtype=torch.long),
            n_options=len(opts),
            truncated=truncated,
        )


class DecidaEncModel(nn.Module):
    """The decision heads on top of the ModernBERT backbone: a type embedding, an optional small transformer
    head, and a scorer that reads one logit out of each option's <mask> position."""

    def __init__(self, encoder: nn.Module, head_layers: int = 2, n_act: int = 2, cfg: dict[str, Any] | None = None):
        super().__init__()
        self.encoder, self.cfg = encoder, cfg or {}
        d = int(encoder.config.hidden_size)
        layer = nn.TransformerEncoderLayer(d_model=d, nhead=max(1, d // 64), dim_feedforward=4 * d, dropout=0.1, batch_first=True, norm_first=True)
        self.head = nn.TransformerEncoder(layer, head_layers, enable_nested_tensor=False) if head_layers > 0 else None
        self.type_emb = nn.Embedding(3, d)
        self.scorer = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(approximate="none"), nn.Linear(d, 1))
        self.act_head = nn.Sequential(nn.Linear(d + 4, 256), nn.GELU(approximate="none"), nn.Linear(256, n_act))
        self.qtypes_map = {"choice": 0, "score": 1, "noul": 2}

    def forward(self, batch: Batch) -> torch.Tensor:
        h = self.encoder(input_ids=batch.input_ids, attention_mask=batch.attention_mask).last_hidden_state
        h = h + self.type_emb(batch.qtypes)[:, None, :]
        if self.head is not None:
            pad = ~batch.attention_mask.bool()
            for layer in self.head.layers:
                h = layer(h, src_key_padding_mask=pad)
        idx = batch.readout_positions.clamp(min=0)[:, :, None].expand(-1, -1, h.size(-1))
        logits = self.scorer(torch.gather(h, 1, idx)).squeeze(-1).float()
        return logits.masked_fill(~batch.marker_mask, float("-inf"))

    @classmethod
    def from_pretrained(cls, path: str) -> DecidaEncModel:
        weights = load_file(os.path.join(path, "model.safetensors"))
        cfg_path = os.path.join(path, "rl_agent_config.json")
        cfg = json.load(open(cfg_path)) if os.path.exists(cfg_path) else {}
        ecfg = AutoConfig.from_pretrained(os.path.join(path, "encoder"))
        enc = AutoModel.from_config(ecfg, attn_implementation="sdpa")

        act_weight = weights.get("act_head.2.weight", weights.get("core_model.act_head.2.weight"))
        if act_weight is None:
            raise ValueError(f"{path} has no act_head.2.weight: not a Decida or Laya encoder checkpoint")
        n_act, head_layers, new_weights = act_weight.shape[0], 0, {}
        for k, v in weights.items():
            k = k.removeprefix("core_model.")
            new_weights[k] = v
            if k.startswith("head.layers."):
                head_layers = max(head_layers, int(k.split(".")[2]) + 1)

        model = cls(enc, head_layers=head_layers, n_act=n_act, cfg=cfg)
        model.load_state_dict(new_weights, strict=False)
        return model


def load(device: str = "cuda") -> tuple[DecidaEncModel, LayaLayout]:
    """Download `helmo/DecidaBERT-large` from the Hub and build the model plus its prompt layout."""
    from huggingface_hub import snapshot_download

    path = snapshot_download(MODEL_REPO)
    model = DecidaEncModel.from_pretrained(path).to(device).eval()
    tokenizer = AutoTokenizer.from_pretrained(os.path.join(path, "tokenizer") if os.path.isdir(os.path.join(path, "tokenizer")) else path)
    # DecidaBERT-large was trained/prompted at a 2142-token context; the checkpoint's own config can override this.
    layout = LayaLayout(tokenizer, max_len=model.cfg.get("max_len", 2142),
                         head_max_len=model.cfg.get("head_max_len", 1400), opt_max_tokens=model.cfg.get("opt_max_tokens", 160))
    return model, layout


@torch.no_grad()
def predict(model: DecidaEncModel, layout: LayaLayout, state: str, questions: dict[str, Question]) -> dict[str, dict[str, Any]]:
    """Answer every question in one batched forward pass. Returns {qid: {type, ...distribution fields}}."""
    items = [{"qid": qid, "encoded": layout.build(state, q), "qtype": model.qtypes_map[q.primitive.value]} for qid, q in questions.items()]
    if not items:
        return {}

    L = max(len(it["encoded"].input_ids) for it in items)
    kmax = max(it["encoded"].n_options for it in items)
    n = len(items)
    pad_id = layout.tok.pad_token_id or 0
    ids = torch.full((n, L), pad_id, dtype=torch.long)
    att = torch.zeros((n, L), dtype=torch.long)
    mpos = torch.zeros((n, kmax), dtype=torch.long)
    mmask = torch.zeros((n, kmax), dtype=torch.bool)
    qtypes = torch.zeros(n, dtype=torch.long)
    for i, it in enumerate(items):
        enc = it["encoded"]
        ids[i, : len(enc.input_ids)] = enc.input_ids
        att[i, : len(enc.attention_mask)] = enc.attention_mask
        mpos[i, : enc.n_options] = enc.readout_positions
        mmask[i, : enc.n_options] = True
        qtypes[i] = it["qtype"]

    device = next(model.parameters()).device
    batch = Batch(ids.to(device), att.to(device), mpos.to(device), mmask.to(device), qtypes.to(device))
    logits = model.forward(batch).cpu().numpy()

    qtypes_inv = {v: k for k, v in model.qtypes_map.items()}
    temp_array = model.cfg.get("temperature", [1.0, 1.0, 1.0])
    temp_by_opt = model.cfg.get("temperature_by_options", {})

    answers: dict[str, dict[str, Any]] = {}
    for r, it in enumerate(items):
        qid, q, k = it["qid"], questions[it["qid"]], it["encoded"].n_options
        z = logits[r, :k]
        size = "2" if k <= 2 else "3-5" if k <= 5 else "6-10" if k <= 10 else "11+"
        bucket = f"{qtypes_inv[it['qtype']]}:{size}"
        t_scale = temp_by_opt.get(bucket, temp_array[it["qtype"]])
        z = z / max(1e-3, float(t_scale))
        p = np.exp(z - z.max())
        p = p / p.sum()

        if q.primitive == Primitive.CHOICE:
            keys = list(q.criteria.keys())
            probs = {kk: float(v) for kk, v in zip(keys, p)}
            answers[qid] = {"type": "choice", "choice": keys[int(p.argmax())], "probabilities": probs, "confidence": float(max(probs.values()))}
        elif q.primitive == Primitive.SCORE:
            probs = {str(i): float(v) for i, v in enumerate(p)}
            answers[qid] = {"type": "score", "score": float((np.arange(k) * p).sum()),
                             "legend": {str(i): str(c) for i, c in enumerate(q.criteria)}, "probabilities": probs, "confidence": float(max(probs.values()))}
        else:
            noul_prob = float(p[1])
            answers[qid] = {"type": "noul", "noul": noul_prob, "confidence": float(noul_prob if noul_prob > 0.5 else 1 - noul_prob)}
    return answers