"""Decision model: a causal-LM backbone plus a pointer head.

For each question the head scores every option's `</opt>` hidden state against the question's `<decide>` hidden
state (scaled dot product in a small projected space); a softmax over those scores is the answer distribution.
The LM head is never run (frozen, but saved so the checkpoint loads as a normal causal LM). A scalar temperature (fitted after training) divides the scores at inference.
"""
import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForCausalLM


class PointerHead(nn.Module):
    def __init__(self, hidden, dim=256):
        super().__init__()
        self.norm = nn.LayerNorm(hidden)
        self.q = nn.Linear(hidden, dim)
        self.k = nn.Linear(hidden, dim)
        self.dim = dim

    def forward(self, h_decide, h_opts):
        """h_decide: [H]; h_opts: [K, H] -> scores [K]."""
        q = self.q(self.norm(h_decide))
        k = self.k(self.norm(h_opts))
        return k @ q / math.sqrt(self.dim)


class DecisionModel(nn.Module):
    def __init__(self, path, head_dim=256, dtype=torch.bfloat16, attn="sdpa", lora_r=0, marker_ids=None):
        super().__init__()
        lm = AutoModelForCausalLM.from_pretrained(path, dtype=dtype, trust_remote_code=True, attn_implementation=attn)
        lm.lm_head.requires_grad_(False)  # never run; kept only so the saved checkpoint is a complete model
        if lora_r:
            from peft import LoraConfig, get_peft_model
            cfg = LoraConfig(r=lora_r, lora_alpha=2 * lora_r, lora_dropout=0.0, bias="none",
                             target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
                             trainable_token_indices={"embed_tokens": list(marker_ids or [])})
            lm = get_peft_model(lm, cfg)
        self.lm = lm
        self.lora = bool(lora_r)
        self.config = lm.config
        self.head = PointerHead(lm.config.hidden_size, head_dim).to(torch.float32)
        self.register_buffer("temperature", torch.ones(()))

    @property
    def backbone(self):
        base = self.lm.get_base_model() if self.lora else self.lm
        return base.model  # decoder stack only

    def forward(self, batch, return_hidden=False):
        """Returns one score vector per question, in the order of batch['questions'].
        With return_hidden, also returns the stacked `<decide>` hidden states [Q, H] (used by the RL value head)."""
        dev = next(self.parameters()).device
        out = self.backbone(input_ids=batch["input_ids"].to(dev), position_ids=batch["position_ids"].to(dev),
                            attention_mask=batch["mask"].to(dev), use_cache=False)
        h = out.last_hidden_state.float()
        scores = []
        for b, d, ends, _, _ in batch["questions"]:
            scores.append(self.head(h[b, d], h[b, ends]) / self.temperature)
        if return_hidden:
            return scores, torch.stack([h[b, d] for b, d, _, _, _ in batch["questions"]])
        return scores


def loss_fn(scores, questions):
    """Cross-entropy against the (possibly soft) target of each question, averaged over questions."""
    losses = []
    for s, (_, _, _, target, _) in zip(scores, questions):
        t = torch.tensor(target, device=s.device, dtype=s.dtype)
        losses.append(-(t * F.log_softmax(s, -1)).sum())
    return torch.stack(losses).mean()