"""Decoding strategies implemented on top of ``model.forward()``.

``model.generate()`` is never called.  Each strategy drives the forward pass
itself and maintains its own KV cache, which is what keeps the cost linear in
the number of generated tokens instead of quadratic: without a cache every new
token would re-encode the entire prefix.

Prompts in a batch have different lengths, so they are **left**-padded.  That
puts every real last token at index -1, which is what the decoder needs, but
it also means the absolute positions are wrong unless ``position_ids`` is
derived from the attention mask rather than from ``arange`` -- which is what
``_initial_forward`` does below.
"""

from __future__ import annotations

import torch
import torch.nn.functional as F


# ---------------------------------------------------------------------------
# Shared plumbing
# ---------------------------------------------------------------------------
def _initial_forward(model, input_ids, attention_mask):
    """Encode the prompt once and return (last-step logits, cache, position)."""
    pos = attention_mask.long().cumsum(-1) - 1
    pos.clamp_(min=0)
    out = model(input_ids=input_ids, attention_mask=attention_mask,
                position_ids=pos, use_cache=True)
    return out.logits[:, -1, :].float(), out.past_key_values, pos[:, -1]


def _step(model, token, attention_mask, past, pos):
    """One incremental decoding step through the cache."""
    out = model(input_ids=token.unsqueeze(-1), attention_mask=attention_mask,
                position_ids=pos.unsqueeze(-1), past_key_values=past,
                use_cache=True)
    return out.logits[:, -1, :].float(), out.past_key_values


def _grow_mask(attention_mask):
    ones = attention_mask.new_ones((attention_mask.size(0), 1))
    return torch.cat([attention_mask, ones], dim=1)


def _finish(seqs, eos_id, prompt_len):
    """Trim each generated sequence at its first EOS."""
    out = []
    for s in seqs:
        if eos_id is not None and eos_id in s:
            s = s[:s.index(eos_id)]
        out.append(s)
    return out


# ---------------------------------------------------------------------------
# 1. Greedy
# ---------------------------------------------------------------------------
@torch.no_grad()
def greedy_decode(model, input_ids, attention_mask, max_new=64, eos_id=None):
    """Always take the arg-max. Deterministic, and the width-1 limit of beam
    search *only* when no length normalisation is applied."""
    B = input_ids.size(0)
    logits, past, pos = _initial_forward(model, input_ids, attention_mask)
    gen = [[] for _ in range(B)]
    done = torch.zeros(B, dtype=torch.bool, device=input_ids.device)
    scores = torch.zeros(B, device=input_ids.device)

    for _ in range(max_new):
        lp = F.log_softmax(logits, dim=-1)
        nxt = lp.argmax(-1)
        scores += lp.gather(1, nxt.unsqueeze(1)).squeeze(1) * (~done)
        for i, t in enumerate(nxt.tolist()):
            if not done[i]:
                gen[i].append(t)
        if eos_id is not None:
            done |= nxt == eos_id
        if bool(done.all()):
            break
        attention_mask = _grow_mask(attention_mask)
        pos = pos + 1
        logits, past = _step(model, nxt, attention_mask, past, pos)
    return _finish(gen, eos_id, input_ids.size(1)), scores.tolist()


# ---------------------------------------------------------------------------
# 2 & 3. Truncated sampling: top-k and top-p
# ---------------------------------------------------------------------------
def _top_k_filter(logits, k):
    """Keep the k highest logits, mask the rest to -inf."""
    if k is None or k <= 0 or k >= logits.size(-1):
        return logits
    kth = logits.topk(k, dim=-1).values[..., -1, None]
    return logits.masked_fill(logits < kth, float("-inf"))


def _top_p_filter(logits, p):
    """Nucleus filtering: keep the smallest prefix of the sorted distribution
    whose cumulative probability first reaches p.

    The shift below is what guarantees the top-1 token always survives: a
    token is removed only if the mass *strictly before* it already reached p,
    so even a single token with probability > p is kept.
    """
    if p is None or p >= 1.0:
        return logits
    srt, idx = logits.sort(descending=True, dim=-1)
    cum = srt.softmax(-1).cumsum(-1)
    remove = cum - srt.softmax(-1) >= p          # mass strictly before token
    srt = srt.masked_fill(remove, float("-inf"))
    return srt.scatter(-1, idx, srt)


@torch.no_grad()
def sample_decode(model, input_ids, attention_mask, max_new=64, eos_id=None,
                  top_k=None, top_p=None, temperature=1.0, generator=None):
    """Ancestral sampling from a truncated, temperature-scaled distribution."""
    B = input_ids.size(0)
    logits, past, pos = _initial_forward(model, input_ids, attention_mask)
    gen = [[] for _ in range(B)]
    done = torch.zeros(B, dtype=torch.bool, device=input_ids.device)
    scores = torch.zeros(B, device=input_ids.device)

    for _ in range(max_new):
        full_lp = F.log_softmax(logits, dim=-1)        # score under the *model*
        z = logits / max(temperature, 1e-6)
        z = _top_k_filter(z, top_k)
        z = _top_p_filter(z, top_p)
        probs = z.softmax(-1)
        nxt = torch.multinomial(probs, 1, generator=generator).squeeze(1)
        scores += full_lp.gather(1, nxt.unsqueeze(1)).squeeze(1) * (~done)
        for i, t in enumerate(nxt.tolist()):
            if not done[i]:
                gen[i].append(t)
        if eos_id is not None:
            done |= nxt == eos_id
        if bool(done.all()):
            break
        attention_mask = _grow_mask(attention_mask)
        pos = pos + 1
        logits, past = _step(model, nxt, attention_mask, past, pos)
    return _finish(gen, eos_id, input_ids.size(1)), scores.tolist()


# ---------------------------------------------------------------------------
# 4. Beam search
# ---------------------------------------------------------------------------
@torch.no_grad()
def beam_search(model, input_ids, attention_mask, beam_width=4, max_new=64,
                eos_id=None, length_penalty=1.0):
    """Keep the ``beam_width`` highest-scoring prefixes at every step.

    Cost per step is ``O(W * |V|)`` for the scoring plus a ``top-W`` over the
    ``W * |V|`` candidates, and the whole batch is run as one ``B*W`` forward,
    so wall time grows roughly linearly in W (memory does too -- the KV cache
    is W copies of the prompt).  Finished hypotheses are moved aside and their
    beam slot is neutralised with a -inf mask so it cannot be extended again.

    Hypotheses are *selected* by summed log-probability divided by
    ``len ** length_penalty``; without that normalisation beam search
    systematically prefers short outputs, since every extra token adds a
    negative number.  The score *returned* is the raw summed log-probability,
    so that it means the same thing as the score the other strategies return
    and the caller can derive a per-token figure the same way for all of them.
    """
    B, L = input_ids.shape
    dev = input_ids.device
    W, V = beam_width, model.config.vocab_size

    logits, past, pos = _initial_forward(model, input_ids, attention_mask)
    lp = F.log_softmax(logits, dim=-1)                     # (B, V)

    # First expansion: from one prompt to W beams.
    top_lp, top_tok = lp.topk(W, dim=-1)                   # (B, W)
    beam_scores = top_lp.reshape(B * W)                    # (B*W,)
    beams = [[int(t)] for t in top_tok.reshape(-1).tolist()]

    past.batch_repeat_interleave(W)
    attention_mask = attention_mask.repeat_interleave(W, dim=0)
    pos = pos.repeat_interleave(W, dim=0)
    cur = top_tok.reshape(-1)

    # Each entry is (selection score, raw summed logprob, tokens).
    finished = [[] for _ in range(B)]
    alive = torch.ones(B * W, dtype=torch.bool, device=dev)

    for i in range(B * W):
        if eos_id is not None and beams[i][0] == eos_id:
            finished[i // W].append((beam_scores[i].item(),
                                     beam_scores[i].item(), []))
            alive[i] = False

    for step in range(1, max_new):
        if not bool(alive.any()):
            break
        attention_mask = _grow_mask(attention_mask)
        pos = pos + 1
        logits, past = _step(model, cur, attention_mask, past, pos)
        lp = F.log_softmax(logits, dim=-1)                  # (B*W, V)

        # A dead beam must not spawn children; -inf everywhere keeps it from
        # ever winning a top-W slot.
        lp = lp.masked_fill(~alive.unsqueeze(1), float("-inf"))
        cand = (beam_scores.unsqueeze(1) + lp).view(B, W * V)
        top_sc, top_ix = cand.topk(W, dim=-1)               # (B, W)
        src_beam = top_ix // V                              # which beam
        tok = top_ix % V                                    # which token

        flat_src = (torch.arange(B, device=dev).unsqueeze(1) * W
                    + src_beam).reshape(-1)
        past.reorder_cache(flat_src)
        new_beams = [beams[i] + [int(t)]
                     for i, t in zip(flat_src.tolist(), tok.reshape(-1).tolist())]
        beams = new_beams
        beam_scores = top_sc.reshape(-1)
        alive = alive[flat_src].clone()
        cur = tok.reshape(-1)

        if eos_id is not None:
            hit = (cur == eos_id) & alive
            for i in hit.nonzero(as_tuple=True)[0].tolist():
                seq = beams[i][:-1]
                norm = len(seq) ** length_penalty if seq else 1.0
                raw = beam_scores[i].item()
                finished[i // W].append((raw / norm, raw, seq))
                alive[i] = False
                beam_scores[i] = float("-inf")

    # Anything still alive at the cap counts as a (length-normalised) candidate.
    for i in range(B * W):
        if alive[i]:
            seq = beams[i]
            norm = len(seq) ** length_penalty if seq else 1.0
            raw = beam_scores[i].item()
            finished[i // W].append((raw / norm, raw, seq))

    out, out_scores = [], []
    for b in range(B):
        if finished[b]:
            _, raw, seq = max(finished[b], key=lambda t: t[0])
        else:
            raw, seq = 0.0, []
        out.append(seq)
        out_scores.append(raw)          # raw sum, not the selection score
    return out, out_scores