"""SPLADE head over a ModernBERT/LateOn MLM backbone.

Sparse vector = fold(max-pool(top_k-gate(log1p(relu(logits - shift))) * pooling_mask))
where the instruction prefix ("[Q] " / "[D] ") is attended by the backbone but
excluded from pooling. Scores are dot products. Load with:

    model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
    q = model.encode(["a query"], kind="query")                # [N, V] float32
    d = model.encode(["a document"])                           # [N, V] float32
    model.score(q, d)                                          # [Nq, Nd] dot
    model.encode_to_dict(["a query"], kind="query", top_k=20)  # {token: weight}
    model.attribute(["a query"], kind="query")                 # + winning input token
    print(model.render(["a query"], kind="query"))             # terminal bar chart
    print(model.highlight(["a document"]))                     # text, fired words lit
"""

from __future__ import annotations

import sys

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, ModernBertConfig, ModernBertForMaskedLM


class SpladeConfig(ModernBertConfig):
    def __init__(
        self,
        logit_shift: float = 0.0,
        position_top_k: int | None = None,
        vocab_fold: str | None = None,
        query_prefix: str = "",
        document_prefix: str = "",
        query_max_length: int = 128,
        doc_max_length: int = 512,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.logit_shift = logit_shift
        self.position_top_k = position_top_k
        self.vocab_fold = vocab_fold
        self.query_prefix = query_prefix
        self.document_prefix = document_prefix
        self.query_max_length = query_max_length
        self.doc_max_length = doc_max_length


class SpladeModel(ModernBertForMaskedLM):
    config_class = SpladeConfig

    def __init__(self, config: SpladeConfig):
        super().__init__(config)
        # [V] canonical-id map for vocab folding, computed once at export from
        # the tokenizer and stored in the checkpoint (identity when unused).
        self.register_buffer(
            "vocab_fold_index", torch.arange(config.vocab_size), persistent=True
        )
        self._tokenizer = None

    # -- forward path ---------------------------------------------------------

    def _token_weights(
        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
    ) -> torch.Tensor:
        """[B, L] tokens -> [B, L, V] per-position activations (pre-pooling)."""
        cfg = self.config
        logits = super().forward(input_ids=input_ids, attention_mask=attention_mask).logits
        weights = torch.log1p(F.relu(logits - cfg.logit_shift))
        if cfg.position_top_k is not None and cfg.position_top_k < weights.shape[-1]:
            # Keep each position's k largest dims (ties keep more than k).
            cutoff = weights.topk(cfg.position_top_k, dim=-1).values[..., -1:]
            weights = weights * (weights >= cutoff)
        return weights

    def _fold(
        self, sparse: torch.Tensor, source_indices: torch.Tensor | None = None
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Reroute each fold group's mass onto its canonical vocab dim.

        `source_indices` [B, V] (max-pool argmax positions) follows the winning
        group member so attribution keeps pointing at a real input position.
        """
        index = self.vocab_fold_index.unsqueeze(0).expand_as(sparse)
        folded = torch.zeros_like(sparse).scatter_reduce(
            1, index, sparse, reduce="amax", include_self=False
        )
        if source_indices is None:
            return folded, None
        winner = (sparse == folded.gather(1, index)) & (sparse > 0)
        folded_sources = torch.full_like(source_indices, -1).scatter_reduce(
            1, index, torch.where(winner, source_indices, -1), reduce="amax", include_self=False
        )
        return folded, folded_sources.clamp_min_(0)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        pooling_mask: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor:
        """[B, L] tokens -> [B, V] sparse activations (dot-product scoring)."""
        if pooling_mask is None:
            pooling_mask = attention_mask
        weights = self._token_weights(input_ids, attention_mask)
        weights = weights * pooling_mask.unsqueeze(-1).to(weights.dtype)
        sparse = weights.max(dim=1).values
        if self.config.vocab_fold is not None:
            sparse, _ = self._fold(sparse)
        return sparse

    @staticmethod
    def score(queries: torch.Tensor, documents: torch.Tensor) -> torch.Tensor:
        """Dot-product relevance scores: [Nq, V] x [Nd, V] -> [Nq, Nd]."""
        return queries @ documents.T

    # -- tokenization ---------------------------------------------------------

    def _get_tokenizer(self):
        if self._tokenizer is None:
            self._tokenizer = AutoTokenizer.from_pretrained(self.config._name_or_path)
        return self._tokenizer

    def _encode_args(self, kind: str, max_length: int | None) -> tuple[str, int]:
        cfg = self.config
        if kind == "query":
            return cfg.query_prefix, max_length or cfg.query_max_length
        if kind == "document":
            return cfg.document_prefix, max_length or cfg.doc_max_length
        raise ValueError("`kind` must be 'query' or 'document'.")

    def _tokenize(self, texts: list[str], prefix: str, max_length: int):
        """-> (input_ids, attention_mask, pooling_mask, char_offsets)."""
        enc = self._get_tokenizer()(
            [prefix + t for t in texts],
            padding=True,
            truncation=True,
            max_length=max_length,
            return_tensors="pt",
            return_offsets_mapping=True,
            return_special_tokens_mask=True,
        )
        pooling_mask = enc["attention_mask"]
        if prefix:
            # Prefix tokens are attended but dropped from pooling: a token is
            # prefix iff its char span starts inside the prefix string.
            # Specials carry a zero-width (0, 0) span, so guard them.
            in_prefix = (enc["offset_mapping"][..., 0] < len(prefix)) & ~enc[
                "special_tokens_mask"
            ].bool()
            pooling_mask = pooling_mask.masked_fill(in_prefix, 0)
        return enc["input_ids"], enc["attention_mask"], pooling_mask, enc["offset_mapping"]

    def _encode_with_sources(self, texts: list[str], prefix: str, max_length: int):
        """One batch through the model, keeping max-pool source positions.

        -> (input_ids, pooling_mask, char_offsets, sparse [B, V], sources [B, V])
        """
        device = next(self.parameters()).device
        ids, attn, pool, offsets = self._tokenize(texts, prefix, max_length)
        ids, attn, pool = ids.to(device), attn.to(device), pool.to(device)
        weights = self._token_weights(ids, attn)
        weights = weights * pool.unsqueeze(-1).to(weights.dtype)
        pooled = weights.max(dim=1)
        sparse, sources = pooled.values, pooled.indices
        if self.config.vocab_fold is not None:
            sparse, sources = self._fold(sparse, sources)
        return ids, pool, offsets, sparse, sources

    # -- encoding APIs --------------------------------------------------------

    @torch.inference_mode()
    def encode(
        self,
        texts: list[str],
        kind: str = "document",
        batch_size: int = 32,
        max_length: int | None = None,
    ) -> torch.Tensor:
        """Encode raw texts -> [N, V] float32 sparse vectors on CPU.

        `kind` ("query" | "document") selects the instruction prefix and the
        default max length.
        """
        prefix, max_length = self._encode_args(kind, max_length)
        device = next(self.parameters()).device
        rows = []
        for start in range(0, len(texts), batch_size):
            ids, attn, pool, _ = self._tokenize(
                texts[start : start + batch_size], prefix, max_length
            )
            rows.append(
                self(ids.to(device), attn.to(device), pool.to(device)).float().cpu()
            )
        return torch.cat(rows)

    @torch.inference_mode()
    def attribute(
        self,
        texts: list[str],
        kind: str = "document",
        top_k: int | None = 25,
        batch_size: int = 32,
        max_length: int | None = None,
        round_to: int = 4,
    ) -> list[list[dict]]:
        """Encode texts and attribute each output dim to its input subtoken.

        Per text, a weight-sorted list of entries
        `{"token", "weight", "source", "position", "expansion"}`:
        `source`/`position` name the input subtoken whose activation won the
        max for that vocab dim (after folding); `expansion` is True when the
        output term is not the source token's own (folded) dim.
        """
        prefix, max_length = self._encode_args(kind, max_length)
        tokenizer = self._get_tokenizer()
        fold_index = self.vocab_fold_index.cpu()
        results = []
        for start in range(0, len(texts), batch_size):
            ids, _, _, sparse, sources = self._encode_with_sources(
                texts[start : start + batch_size], prefix, max_length
            )
            for row, row_sources, row_ids in zip(
                sparse.float().cpu(), sources.cpu(), ids.cpu()
            ):
                dims = torch.nonzero(row, as_tuple=False).flatten()
                order = torch.argsort(row[dims], descending=True)[:top_k]
                entries = []
                for dim in dims[order].tolist():
                    pos = int(row_sources[dim])
                    src_id = int(row_ids[pos])
                    entries.append(
                        {
                            "token": tokenizer.convert_ids_to_tokens(dim),
                            "weight": round(float(row[dim]), round_to),
                            "source": tokenizer.convert_ids_to_tokens(src_id),
                            "position": pos,
                            "expansion": int(fold_index[src_id]) != dim,
                        }
                    )
                results.append(entries)
        return results

    def encode_to_dict(
        self, texts: list[str], kind: str = "document", top_k: int | None = None, **kwargs
    ) -> list[dict[str, float]]:
        """Encode texts -> {token: weight} dicts sorted by descending weight."""
        return [
            {e["token"]: e["weight"] for e in entries}
            for entries in self.attribute(texts, kind=kind, top_k=top_k, **kwargs)
        ]

    # -- terminal displays ----------------------------------------------------

    _FADE = "▓▒░"
    _HEAT = (196, 202, 208, 214, 220, 190, 108, 66, 60, 241)  # ANSI-256, hot -> cold

    def _heat(self, ratio: float) -> int:
        return self._HEAT[min(int((1 - ratio) * len(self._HEAT)), len(self._HEAT) - 1)]

    @staticmethod
    def _display(token: str) -> str:
        """Strip the Ġ word marker; dot-prefix continuation pieces."""
        if token.startswith("Ġ"):
            return token[1:]
        if token.startswith("["):  # specials: [CLS], [SEP], [Q], [D]
            return token
        return "·" + token

    def render(
        self,
        texts: list[str],
        kind: str = "document",
        top_k: int | None = 25,
        width: int = 36,
        color: bool | None = None,
        **attribute_kwargs,
    ) -> str:
        """Terminal bar chart of the sparse expansions, with attributions.

        One block per text: an `L0` line with the total number of active dims,
        then bars proportional to weight (peak-normalized), each line ending
        with the input subtoken that produced the dim and `<exp>` for pure
        expansions. `color=None` auto-detects a TTY.
        """
        if color is None:
            color = sys.stdout.isatty()
        blocks = []
        for text, entries in zip(
            texts, self.attribute(texts, kind=kind, top_k=None, **attribute_kwargs)
        ):
            shown = text if len(text) <= 70 else text[:67] + "..."
            if not entries:
                blocks.append(f"{kind} · {shown}\n  (empty vector)")
                continue
            total = len(entries)
            entries = entries[:top_k]
            l0 = f"L0 = {total} active dims"
            if len(entries) < total:
                l0 += f" (showing {len(entries)})"
            peak = entries[0]["weight"]
            name_width = max(len(self._display(e["token"])) for e in entries)
            lines = [f"{kind} · {shown}", f"\033[2m{l0}\033[0m" if color else l0]
            for e in entries:
                ratio = e["weight"] / peak
                cells = max(1, round(ratio * width))
                bar = ("█" * cells)[:-3] + self._FADE if cells > 3 else self._FADE[3 - cells :]
                pad = " " * (width - cells)
                token = self._display(e["token"]).rjust(name_width)
                attrib = f"<- {self._display(e['source'])}@{e['position']}"
                if e["expansion"]:
                    attrib += " <exp>"
                if color:
                    heat = self._heat(ratio)
                    token = f"\033[38;5;{heat}m{token}\033[0m"
                    bar = f"\033[38;5;{heat}m{bar}\033[0m"
                    attrib = f"\033[2m{attrib}\033[0m"
                lines.append(f"{token} {bar}{pad} {e['weight']:>6.2f}  {attrib}")
            blocks.append("\n".join(lines))
        return "\n\n".join(blocks)

    @torch.inference_mode()
    def highlight(
        self,
        texts: list[str],
        kind: str = "document",
        reduce: str = "sum",
        color: bool | None = None,
        batch_size: int = 32,
        max_length: int | None = None,
    ) -> str:
        """Render each text with its firing words lit up.

        A word fires when one of its subtokens wins the max for at least one
        output dim; intensity is the total mass it contributes (`reduce="sum"`,
        default) or the largest single weight it wins (`reduce="max"`). TTY:
        reverse-video heat colors. Plain: tiered markers `⟦strong⟧ «mid» ‹weak›`.
        """
        if reduce not in ("sum", "max"):
            raise ValueError("`reduce` must be 'sum' or 'max'.")
        if color is None:
            color = sys.stdout.isatty()
        prefix, max_length = self._encode_args(kind, max_length)
        blocks = []
        for start in range(0, len(texts), batch_size):
            batch = texts[start : start + batch_size]
            ids, pool, offsets, sparse, sources = self._encode_with_sources(
                batch, prefix, max_length
            )
            # [B, L] per-position intensity, reduced over the dims each position
            # won. Zero-weight dims carry a clamped position 0 but contribute 0,
            # so they can't corrupt the reduction.
            intensity = torch.zeros(
                ids.shape, dtype=sparse.dtype, device=sparse.device
            ).scatter_reduce(
                1, sources, sparse, reduce="sum" if reduce == "sum" else "amax", include_self=False
            )
            for text, row_int, row_off, row_pool in zip(
                batch, intensity.float().cpu(), offsets, pool.cpu()
            ):
                blocks.append(
                    self._paint(text, row_int, row_off, row_pool, len(prefix), color, reduce)
                )
        return "\n\n".join(blocks)

    def _paint(self, text, intensity, offsets, pooling_mask, prefix_len, color, reduce) -> str:
        """Wrap fired char spans of `text` in intensity markers."""
        # Byte-BPE offsets include the word's leading space: trim it, so merging
        # only fuses glued subtokens of the same word (one span per word).
        spans: list[list] = []
        for pos in range(len(offsets)):
            w = intensity[pos].item()
            if w <= 0 or pooling_mask[pos] == 0:
                continue
            s, e = int(offsets[pos][0]) - prefix_len, int(offsets[pos][1]) - prefix_len
            s = max(s, 0)
            while s < e and text[s].isspace():
                s += 1
            if e <= s:  # zero-width specials / whitespace-only
                continue
            if spans and s == spans[-1][1]:
                spans[-1][1] = e
                spans[-1][2] = spans[-1][2] + w if reduce == "sum" else max(spans[-1][2], w)
            else:
                spans.append([s, e, w])
        if not spans:
            return text
        peak = max(w for _, _, w in spans)  # after merging, so summed words stay <= 1
        out, cursor = [], 0
        for s, e, w in spans:
            ratio = w / peak
            out.append(text[cursor:s])
            if color:
                # Reverse video with heat foreground = heat-colored highlighter.
                out.append(f"\033[7;38;5;{self._heat(ratio)}m{text[s:e]}\033[0m")
            else:
                marks = "⟦⟧" if ratio > 0.66 else "«»" if ratio > 0.33 else "‹›"
                out.append(f"{marks[0]}{text[s:e]}{marks[1]}")
            cursor = e
        out.append(text[cursor:])
        return "".join(out)