"""Prompt lookup decoding (PLD): drafts copied from the context (prompt + output).

At each cycle, the longest suffix of the context of length n_max..n_min that occurred
earlier is matched to its MOST RECENT previous occurrence, and the m tokens that followed
it are proposed as drafts. The drafts are verified like any other (exact, lossless).
"""


class NgramIndex:
    """Incremental index: for n in [n_min, n_max], the last END position of every n-gram,
    excluding the gram that ends at the current last token (so a match is always earlier)."""

    def __init__(self, n_max: int = 4, n_min: int = 1):
        self.n_max, self.n_min = n_max, n_min
        self.seq: list[int] = []
        self.last = {n: {} for n in range(n_min, n_max + 1)}

    def extend(self, toks) -> None:
        for t in toks:
            self.seq.append(t)
            e = len(self.seq) - 2  # index grams ending one before the new last token
            if e < 0:
                continue
            for n in range(self.n_min, self.n_max + 1):
                if e - n + 1 >= 0:
                    self.last[n][tuple(self.seq[e - n + 1:e + 1])] = e

    def propose(self, m: int) -> list[int]:
        """Up to m copied tokens; self.last_n = length of the matched suffix (0 if none)."""
        s = self.seq
        self.last_n = 0
        for n in range(self.n_max, self.n_min - 1, -1):
            if len(s) < n:
                continue
            e = self.last[n].get(tuple(s[-n:]))
            if e is not None:
                self.last_n = n
                return s[e + 1:e + 1 + m]
        return []