"""The decision model around a Hugging Face causal LM: two-way reading and a code-row readout.

**Two-way reading.** With ``attention="two_way"`` the full-attention layers see the whole prompt
in both directions, while the linear-attention layers (Qwen3.5's Gated DeltaNet) and any
sliding-window layers keep the model's own masks.

**Readout.** The hidden state at each question's readout position is scored against one row per
option code: the LM head's own rows of the code tokens (``readout="lm_head"``), or a separate
``[n_codes, d]`` matrix (``readout="head"``). The logits are normalised over the question's
options only.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Sequence

import torch
from torch import nn

from glance_runtime.layout import Example


def _core(lm: nn.Module) -> nn.Module:
    """The decoder stack (the module that builds attention masks); PEFT wrappers are unwrapped."""
    if hasattr(lm, "get_base_model"):
        lm = lm.get_base_model()
    for attr in ("model", "language_model", "transformer"):
        m = getattr(lm, attr, None)
        if isinstance(m, nn.Module):
            inner = getattr(m, "language_model", None)
            return inner if isinstance(inner, nn.Module) else m
    return lm


def _multimodal(lm: nn.Module) -> nn.Module | None:
    """The module holding the vision encoder and M-RoPE indexing (Qwen3.5's ``model``), if any."""
    if hasattr(lm, "get_base_model"):
        lm = lm.get_base_model()
    for m in (getattr(lm, "model", None), lm):
        if isinstance(m, nn.Module) and hasattr(m, "get_image_features") and hasattr(m, "get_rope_index") \
                and hasattr(m, "get_placeholder_mask"):
            return m
    return None


def _text_config(lm: nn.Module):
    cfg = lm.config
    return getattr(cfg, "text_config", None) or cfg


def two_way_masks(lm: nn.Module, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor) -> Any:
    """Bidirectional masks for full-attention layers; the model's own masks for the others."""
    from transformers import masking_utils as mu

    cfg = _core(lm).config
    kw = {"config": cfg, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, "past_key_values": None}
    full = mu.create_bidirectional_mask(**kw, allow_is_bidirectional_skip=False)
    layer_types = getattr(cfg, "layer_types", None)
    if not layer_types:
        return full
    masks: dict[str, Any] = {}
    pos = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device)[None].expand(inputs_embeds.shape[0], -1)
    for t in set(layer_types):
        if t == "full_attention":
            masks[t] = full
        elif t == "linear_attention":
            masks[t] = mu.create_recurrent_attention_mask(**kw)
        elif t == "sliding_attention":
            masks[t] = mu.create_sliding_window_causal_mask(**kw, position_ids=pos)
        else:
            masks[t] = mu.create_causal_mask(**kw, position_ids=pos)
    return masks


def varlen_masks(lm: nn.Module, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, causal: bool) -> dict:
    """Masks for ``attn_impl="varlen"``: packing indices for full-attention layers (two-way or causal),
    the model's own masks for the others."""
    from transformers import masking_utils as mu

    from glance_runtime.attention import VarlenMask

    cfg = _core(lm).config
    kw = {"config": cfg, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, "past_key_values": None}
    pos = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device)[None].expand(inputs_embeds.shape[0], -1)
    masks: dict[str, Any] = {}
    for t in set(getattr(cfg, "layer_types", None) or ["full_attention"]):
        if t == "full_attention":
            masks[t] = VarlenMask.from_padding(attention_mask, causal=causal)
        elif t == "linear_attention":
            masks[t] = mu.create_recurrent_attention_mask(**kw)
        else:
            raise ValueError(f"attn_impl varlen does not handle {t!r} layers")
    return masks


def set_full_attention_causal(lm: nn.Module, causal: bool) -> int:
    """Flip ``is_causal`` on full-attention modules (SDPA and flash paths read it). Returns how many."""
    n = 0
    for name, mod in lm.named_modules():
        cls = type(mod).__name__
        if cls.endswith("Attention") and "Vision" not in cls and hasattr(mod, "is_causal") and \
                getattr(mod, "layer_type", "full_attention") != "sliding_attention":
            mod.is_causal = causal
            n += 1
    return n


@dataclass
class Batch:
    input_ids: torch.Tensor  # (B, T)
    attention_mask: torch.Tensor  # (B, T)
    readout_row: torch.Tensor  # (R,) batch row of each readout
    readout_pos: torch.Tensor  # (R,) position of each readout
    option_code: torch.Tensor  # (N,) code index of each (readout, option) pair, flattened
    option_readout: torch.Tensor  # (N,) readout index of each pair
    sizes: list[int]  # options per readout
    readouts: list  # layout.Readout objects, in order
    pixel_values: torch.Tensor | None = None  # every image's patches, in example then image order
    image_grid_thw: torch.Tensor | None = None  # (n_images, 3)

    def to(self, device: Any) -> "Batch":
        """Move the tensors the forward pass reads (indices stay where ``forward`` needs them)."""
        self.input_ids = self.input_ids.to(device)
        self.attention_mask = self.attention_mask.to(device)
        if self.pixel_values is not None:
            self.pixel_values = self.pixel_values.to(device)
            self.image_grid_thw = self.image_grid_thw.to(device)
        return self


def collate(examples: Sequence[Example], code_index: dict[str, int], pad_id: int = 0) -> Batch:
    t = max(len(e.input_ids) for e in examples)
    ids = torch.full((len(examples), t), pad_id, dtype=torch.long)
    mask = torch.zeros((len(examples), t), dtype=torch.long)
    rows, poss, codes, owners, sizes, ros = [], [], [], [], [], []
    for b, e in enumerate(examples):
        ids[b, : len(e.input_ids)] = torch.tensor(e.input_ids)
        mask[b, : len(e.input_ids)] = 1
        for ro in e.readouts:
            r = len(ros)
            rows.append(b)
            poss.append(ro.position)
            codes.extend(code_index[c] for c in ro.codes)
            owners.extend([r] * len(ro.codes))
            sizes.append(len(ro.codes))
            ros.append(ro)
    vision = [v for e in examples for v in e.vision]
    pixels = torch.cat([v["pixel_values"] for v in vision]) if vision else None
    grids = torch.cat([v["image_grid_thw"] for v in vision]) if vision else None
    return Batch(ids, mask, torch.tensor(rows), torch.tensor(poss), torch.tensor(codes), torch.tensor(owners), sizes, ros,
                 pixels, grids)


class DecisionModel(nn.Module):
    def __init__(self, lm: nn.Module, code_token_ids: Sequence[int], codes: Sequence[str], *,
                 attention: str = "two_way", readout: str = "lm_head", attn_impl: str = "sdpa"):
        super().__init__()
        if attention not in ("two_way", "causal"):
            raise ValueError("attention must be two_way or causal")
        if attn_impl not in ("sdpa", "varlen"):
            raise ValueError("attn_impl must be sdpa or varlen")
        if readout not in ("lm_head", "head"):
            raise ValueError("readout must be lm_head or head")
        self.lm = lm
        self.attention = attention
        self.readout = readout
        self.codes = list(codes)
        self.code_index = {c: i for i, c in enumerate(self.codes)}
        self.register_buffer("code_token_ids", torch.tensor(list(code_token_ids), dtype=torch.long), persistent=False)
        rows = self.lm_head_weight()[self.code_token_ids].detach().clone()
        self.head = nn.Parameter(rows) if readout == "head" else None
        self.image_encoder: Any = None  # a glance_runtime.vision.ImageEncoder for vision bases
        set_full_attention_causal(lm, attention == "causal")
        self.attn_impl = attn_impl
        if attn_impl == "varlen":  # unpadded full attention on the text model only (glance_runtime.attention)
            from glance_runtime import attention as att

            att.use(_text_config(lm))

    def lm_head_weight(self) -> torch.Tensor:
        return self.lm.get_output_embeddings().weight

    def code_rows(self) -> torch.Tensor:
        return self.head if self.head is not None else self.lm_head_weight()[self.code_token_ids]

    def hidden(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, pixel_values: torch.Tensor | None = None,
               image_grid_thw: torch.Tensor | None = None) -> torch.Tensor:
        core = _core(self.lm)
        embeds = core.get_input_embeddings()(input_ids)
        position_ids = None
        if pixel_values is not None:
            mm = _multimodal(self.lm)
            if mm is None:
                raise ValueError("pixel_values given to a model without a vision encoder")
            feats = mm.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output
            feats = torch.cat(list(feats), 0).to(embeds.device, embeds.dtype)
            image_mask, _ = mm.get_placeholder_mask(input_ids, inputs_embeds=embeds, image_features=feats)
            embeds = embeds.masked_scatter(image_mask, feats)
            types = (input_ids == mm.config.image_token_id).int()
            position_ids, _ = mm.get_rope_index(input_ids, mm_token_type_ids=types, image_grid_thw=image_grid_thw,
                                                attention_mask=attention_mask)
        mask: Any = attention_mask
        if self.attn_impl == "varlen":
            mask = varlen_masks(self.lm, embeds, attention_mask, causal=self.attention == "causal")
        elif self.attention == "two_way":
            mask = two_way_masks(self.lm, embeds, attention_mask)
        out = core(inputs_embeds=embeds, attention_mask=mask, position_ids=position_ids, use_cache=False)
        return out.last_hidden_state

    def forward(self, batch: Batch) -> list[torch.Tensor]:
        """One logit vector per readout, over that question's options in display order."""
        h = self.hidden(batch.input_ids, batch.attention_mask, batch.pixel_values, batch.image_grid_thw)
        hr = h[batch.readout_row, batch.readout_pos]  # (R, d)
        w = self.code_rows()[batch.option_code]  # (N, d)
        flat = (hr[batch.option_readout] * w.to(hr.dtype)).sum(-1)
        return list(torch.split(flat, batch.sizes))