"""Backbone + one scoring head. No lm_head, no generation.

A decision is one forward pass: render state, question and candidates; take the
backbone's final hidden state at each candidate's last token; a small MLP turns
each into a logit; softmax over the candidates. The backbone's pretraining is
where the world knowledge comes from, LoRA adapts it, and the head is the only
part that is new.

Calibration is a temperature per question type, fitted on held-out data after
training (see evaluate.py) and stored beside the weights. It changes no
answer, only how sure the model says it is.
"""

from __future__ import annotations

import json
import math
import os

import torch
from torch import nn

from .render import Renderer
from .schema import Decision


class Head(nn.Module):
    def __init__(self, hidden: int):
        super().__init__()
        mid = max(64, hidden // 4)
        self.net = nn.Sequential(nn.LayerNorm(hidden), nn.Linear(hidden, mid),
                                 nn.GELU(), nn.Linear(mid, 1))

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        return self.net(h).squeeze(-1)


def lora_targets(model) -> list[str]:
    """Attention and MLP projections, by whatever names this family uses."""
    names = {n.rsplit(".", 1)[-1] for n, m in model.named_modules()
             if isinstance(m, nn.Linear) or type(m).__name__ == "Conv1D"}
    llama = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
    got = [n for n in llama if n in names]
    if got:
        return got
    return [n for n in ("c_attn", "c_proj", "c_fc") if n in names]   # gpt2


class DecisionModel(nn.Module):
    def __init__(self, backbone, tok, config: dict):
        super().__init__()
        self.backbone = backbone
        self.tok = tok
        self.config = config
        hidden = backbone.config.hidden_size
        self.head = Head(hidden)
        self.renderer = Renderer(tok, config.get("max_state_tokens", 384),
                                 config.get("max_option_tokens", 32))
        self.temperature = {t: 1.0 for t in ("choice", "score", "null")}

    # ---- construction ---------------------------------------------------- #
    @classmethod
    def create(cls, base: str, lora_rank: int = 16, lora_alpha: int = 32,
               dtype: torch.dtype = torch.float32, gradient_checkpointing: bool = False,
               **cfg) -> "DecisionModel":
        from peft import LoraConfig, get_peft_model
        from transformers import AutoModel, AutoTokenizer

        tok = AutoTokenizer.from_pretrained(base)
        if tok.pad_token_id is None:
            tok.pad_token = tok.eos_token
        bb = AutoModel.from_pretrained(base, dtype=dtype)
        if gradient_checkpointing:
            bb.gradient_checkpointing_enable(
                gradient_checkpointing_kwargs={"use_reentrant": False})
            bb.enable_input_require_grads()
        bb.config.use_cache = False
        targets = lora_targets(bb)
        bb = get_peft_model(bb, LoraConfig(
            r=lora_rank, lora_alpha=lora_alpha, lora_dropout=0.05,
            target_modules=targets, bias="none"))
        config = {"base": base, "lora_rank": lora_rank, "lora_alpha": lora_alpha,
                  "lora_targets": targets, **cfg}
        return cls(bb, tok, config)

    def save(self, path: str) -> None:
        os.makedirs(path, exist_ok=True)
        self.backbone.save_pretrained(os.path.join(path, "adapter"))
        self.tok.save_pretrained(os.path.join(path, "tokenizer"))
        # safetensors, not a pickle: loading a .pt can run code, and a model
        # meant to be downloaded should not ask anyone to trust that.
        from safetensors.torch import save_file
        save_file({k: v.detach().cpu().contiguous() for k, v in self.head.state_dict().items()},
                  os.path.join(path, "head.safetensors"))
        with open(os.path.join(path, "decisionx.json"), "w") as fh:
            json.dump({"config": self.config, "temperature": self.temperature},
                      fh, indent=2)

    @classmethod
    def load(cls, path: str, dtype: torch.dtype = torch.float32,
             device: str | None = None) -> "DecisionModel":
        from peft import PeftModel
        from transformers import AutoModel, AutoTokenizer

        meta = json.load(open(os.path.join(path, "decisionx.json")))
        cfg = meta["config"]
        tok = AutoTokenizer.from_pretrained(os.path.join(path, "tokenizer"))
        bb = AutoModel.from_pretrained(cfg["base"], dtype=dtype)
        bb = PeftModel.from_pretrained(bb, os.path.join(path, "adapter"))
        m = cls(bb, tok, cfg)
        st = os.path.join(path, "head.safetensors")
        if os.path.exists(st):
            from safetensors.torch import load_file
            m.head.load_state_dict(load_file(st))
        else:                                        # saved before v2's export
            m.head.load_state_dict(torch.load(os.path.join(path, "head.pt"),
                                              map_location="cpu", weights_only=True))
        m.temperature.update(meta.get("temperature") or {})
        m.eval()
        if device:
            m.to(device)
        return m

    # ---- the forward pass ------------------------------------------------ #
    def forward(self, input_ids, attention_mask, option_pos) -> torch.Tensor:
        """(B, K) logits; padded candidate slots (option_pos == -1) are -inf."""
        out = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        h = out.last_hidden_state                                   # (B, L, H)
        idx = option_pos.clamp(min=0)
        picked = torch.gather(h, 1, idx.unsqueeze(-1).expand(-1, -1, h.size(-1)))
        logits = self.head(picked.float())                          # (B, K)
        return logits.masked_fill(option_pos < 0, float("-inf"))

    # ---- the API ----------------------------------------------------------- #
    @torch.no_grad()
    def decide(self, type: str, question: str, state: str = "",
               options: list[str] | None = None, scale: dict | None = None,
               calibration=None) -> dict:
        """One decision. Returns the answer, a probability per candidate, and
        a confidence; the shape of the answer depends on the type. Pass the
        result of `calibrate()` to apply a per-task calibration."""
        d = Decision(type=type, state=state, question=question,
                     options=list(options or []), scale=scale).validate()
        return self.decide_many([d], calibration)[0]

    @torch.no_grad()
    def raw_logits(self, decisions: list[Decision]) -> list[torch.Tensor]:
        """Uncalibrated logits over each decision's candidates."""
        from .data.collate import collate
        dev = next(self.parameters()).device
        batch = collate([self.renderer.render(d) for d in decisions],
                        self.tok.pad_token_id)
        logits = self(batch["input_ids"].to(dev), batch["attention_mask"].to(dev),
                      batch["option_pos"].to(dev)).float().cpu()
        return [logits[i, :len(d.candidates())] for i, d in enumerate(decisions)]

    @torch.no_grad()
    def decide_many(self, decisions: list[Decision], calibration=None) -> list[dict]:
        return [self._answer(d, lg, calibration)
                for d, lg in zip(decisions, self.raw_logits(decisions))]

    @torch.no_grad()
    def calibrate(self, examples: list[Decision], strength: float | None = None,
                  batch_size: int = 16):
        """Fit a calibration for one task from labelled examples (20-50 is
        plenty). Every example must be the same type; for a choice task they
        should share an option set, in any order."""
        from .calibration import fit, fit_auto
        if not examples or any(d.target() is None for d in examples):
            raise ValueError("calibration needs labelled examples")
        types = {d.type for d in examples}
        if len(types) != 1:
            raise ValueError(f"one task at a time; got types {sorted(types)}")
        rows = []
        for i in range(0, len(examples), batch_size):
            chunk = examples[i:i + batch_size]
            for d, lg in zip(chunk, self.raw_logits(chunk)):
                rows.append((lg.tolist(), d.candidates(), d.target()))
        t = types.pop()
        T0 = self.temperature.get(t, 1.0)
        if strength is None:         # chosen by cross-validation on these examples
            return fit_auto(rows, t, T0, examples[0].question)
        return fit(rows, t, T0, strength, examples[0].question)

    def _answer(self, d: Decision, logits: torch.Tensor, calibration=None) -> dict:
        if calibration is not None:
            p = calibration.apply(d, logits)
        else:
            p = torch.softmax(logits / self.temperature.get(d.type, 1.0), -1)
        probs = [round(float(x), 6) for x in p]
        best = int(p.argmax())
        if d.type == "null":
            return {"type": "null", "probability": probs[0], "answer": probs[0] >= 0.5}
        if d.type == "score":
            vals = d.scale_values()
            mean = sum(v * q for v, q in zip(vals, probs))
            sd = math.sqrt(sum(q * (v - mean) ** 2 for v, q in zip(vals, probs)))
            return {"type": "score", "score": round(mean, 4), "mode": vals[best],
                    "spread": round(sd, 4), "distribution": dict(zip(vals, probs))}
        return {"type": "choice", "answer": d.options[best], "index": best,
                "confidence": probs[best], "probabilities": dict(zip(d.options, probs))}