"""Model wrapper used at inference: PEFT-wrapped causal LM -> option logits through the restricted LM head (+ residual)."""
import os
import sys
from typing import Any, Dict, List, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import answer_slot as AS  # noqa: E402


class DensePipeline(nn.Module):
    """PEFT-wrapped causal LM -> option logits through the restricted LM head (+ residual)."""

    def __init__(self, peft_model: Any, cand_ids: List[int], softcap: Optional[float] = None):
        super().__init__()
        self.peft = peft_model
        base = peft_model.get_base_model() if hasattr(peft_model, "get_base_model") else peft_model   # merged models have no PEFT wrapper
        self.backbone = base.model
        hidden = base.config.get_text_config().hidden_size
        w = base.lm_head.weight
        self.register_buffer("cand_w", w[torch.tensor(cand_ids, device=w.device)].detach().clone(), persistent=False)
        self.head = AS.SlotResidual(hidden, len(cand_ids))
        self.softcap = softcap

    def readout(self, hs, slot_of_opt, pos_idx, neg_idx, opt_mask):
        """Option logits from answer-slot hidden states hs (B, S, H)."""
        lin = F.linear(hs, self.cand_w).float()
        if self.softcap:
            lin = self.softcap * torch.tanh(lin / self.softcap)
        z = lin + self.head(hs).float()
        zs = z.gather(1, slot_of_opt.unsqueeze(-1).expand(-1, -1, z.size(-1)))
        pos = zs.gather(2, pos_idx.unsqueeze(-1)).squeeze(-1)
        neg = zs.gather(2, neg_idx.clamp(min=0).unsqueeze(-1)).squeeze(-1)
        return (pos - neg * (neg_idx >= 0)).masked_fill(~opt_mask, -1e4)

    def forward(self, input_ids, attention_mask, slot_pos, slot_of_opt, pos_idx, neg_idx, opt_mask, return_pooled=False):
        h = self.backbone(input_ids=input_ids, attention_mask=attention_mask, use_cache=False).last_hidden_state
        rows = torch.arange(h.size(0), device=h.device).unsqueeze(1)
        hs = h[rows, slot_pos]
        logits = self.readout(hs, slot_of_opt, pos_idx, neg_idx, opt_mask)
        if return_pooled:   # mean answer-slot hidden state (temperature head input)
            valid = (torch.arange(slot_pos.size(1), device=h.device)[None, :] < (slot_pos > 0).sum(1, keepdim=True)).unsqueeze(-1)
            return logits, (hs.float() * valid).sum(1) / valid.sum(1).clamp(min=1)
        return logits




def load_model(path: str, quant: str):
    from transformers import AutoModelForCausalLM, BitsAndBytesConfig
    kw: Dict[str, Any] = {"dtype": torch.bfloat16, "device_map": "cuda"}
    if quant == "nf4":
        kw["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
    return AutoModelForCausalLM.from_pretrained(path, **kw)