"""Decoding strategies, worked as Chapter 18 works them.

Four things, all reproducing the chapter's tables:

    greedy against beam   where greedy loses, and what beam costs
    temperature           one knob, three behaviours
    top-k and top-p       why a fixed k fails on some distributions
    the whole picture     the same distribution under every strategy

    python3 decoding.py            # all four
    python3 decoding.py --beam     # greedy against beam search
    python3 decoding.py --temp     # the temperature sweep
    python3 decoding.py --truncate # top-k against top-p
    python3 decoding.py --compare  # every strategy, side by side

Install:  nothing, the Python standard library is enough
"""

import argparse
import math

# A tiny language model given as an explicit tree of next-token probabilities.
# The key is the prefix; the value maps next token to probability.
TREE = {
    (): {"the": 0.40, "a": 0.35, "one": 0.25},
    ("the",): {"cat": 0.55, "dog": 0.45},
    ("a",): {"cat": 0.10, "bird": 0.90},
    ("one",): {"cat": 0.50, "dog": 0.50},
    ("the", "cat"): {"sat": 0.60, "ran": 0.40},
    ("the", "dog"): {"sat": 0.50, "ran": 0.50},
    ("a", "cat"): {"sat": 0.50, "ran": 0.50},
    ("a", "bird"): {"sang": 0.95, "flew": 0.05},
    ("one", "cat"): {"sat": 0.50, "ran": 0.50},
    ("one", "dog"): {"sat": 0.50, "ran": 0.50},
}

# A separate distribution for the truncation demonstrations. Two shapes: one
# where the model is confident and one where it is not.
PEAKED = {"the": 0.85, "a": 0.07, "one": 0.04, "some": 0.02,
          "any": 0.01, "each": 0.005, "every": 0.005}
FLAT = {"red": 0.14, "blue": 0.13, "green": 0.13, "black": 0.12,
        "white": 0.12, "grey": 0.12, "brown": 0.12, "pink": 0.12}


def greedy(depth=3):
    prefix, logp, trace = (), 0.0, []
    for _ in range(depth):
        dist = TREE.get(prefix)
        if not dist:
            break
        w = max(dist, key=dist.get)
        trace.append((prefix, w, dist[w], dict(dist)))
        logp += math.log(dist[w])
        prefix = prefix + (w,)
    return prefix, math.exp(logp), trace


def all_sequences(depth=3):
    out = []

    def walk(prefix, p):
        dist = TREE.get(prefix)
        if not dist or len(prefix) == depth:
            out.append((prefix, p))
            return
        for w, q in dist.items():
            walk(prefix + (w,), p * q)
    walk((), 1.0)
    return sorted(out, key=lambda x: -x[1])


def beam(width, depth=3):
    beams = [((), 0.0)]
    for _ in range(depth):
        cand = []
        for prefix, lp in beams:
            dist = TREE.get(prefix)
            if not dist:
                cand.append((prefix, lp))
                continue
            for w, q in dist.items():
                cand.append((prefix + (w,), lp + math.log(q)))
        cand.sort(key=lambda x: -x[1])
        beams = cand[:width]
    return [(p, math.exp(lp)) for p, lp in beams]


# --------------------------------------------------------------- beam search

def show_beam():
    print("\nSTEP 1  greedy takes the best word and loses the best sentence\n")
    seq, p, trace = greedy()
    print("  GREEDY. At each step take the highest probability token.\n")
    print(f"  {'prefix':<18}{'choices':<34}{'taken':<8}{'p':>7}")
    print("  " + "-" * 68)
    for prefix, w, q, dist in trace:
        pre = " ".join(prefix) if prefix else "(start)"
        ch = ", ".join(f"{k} {v:.2f}" for k, v in dist.items())
        print(f"  {pre:<18}{ch:<34}{w:<8}{q:>7.2f}")
    print(f"\n  greedy output: '{' '.join(seq)}'   probability {p:.4f}\n")

    ranked = all_sequences()
    print("  Now every three-word sequence, ranked by probability:\n")
    print(f"  {'rank':>5}  {'sequence':<24}{'probability':>13}")
    print("  " + "-" * 46)
    for i, (s, q) in enumerate(ranked[:6], start=1):
        mark = "   <- greedy found this" if s == seq else ""
        print(f"  {i:>5}  {' '.join(s):<24}{q:>13.4f}{mark}")

    best = ranked[0]
    print(f"\n  The best sequence is '{' '.join(best[0])}' at {best[1]:.4f}.")
    print(f"  Greedy returned {p:.4f}, which is"
          f" {(best[1] / p - 1) * 100:.0f} per cent worse.")
    print(f"\n  Greedy went wrong at the very first step. It took 'the' at")
    print(f"  0.40 over 'a' at 0.35, and never saw that 'a' leads to 'bird'")
    print(f"  at 0.90 and then 'sang' at 0.95.")
    print(f"\n  One high-probability token can sit in front of a low-")
    print(f"  probability continuation. Greedy cannot know that in advance.")

    print(f"\n\n  BEAM SEARCH keeps the best w partial sequences at each step.\n")
    print(f"  {'width':>6}  {'best sequence found':<26}{'probability':>13}"
          f"{'optimal?':>10}")
    print("  " + "-" * 58)
    for w in (1, 2, 3, 5):
        b = beam(w)
        top = b[0]
        ok = "yes" if abs(top[1] - best[1]) < 1e-12 else "no"
        print(f"  {w:>6}  {' '.join(top[0]):<26}{top[1]:>13.4f}{ok:>10}")

    print(f"\n  Width 1 is greedy by definition. Width 2 already finds the")
    print(f"  optimum here, and wider beams cost more for nothing.")
    print(f"\n  Beam search is not guaranteed to find the best sequence. It")
    print(f"  searches more of the space than greedy and less than all of it,")
    print(f"  which is the only honest description of it.")


# --------------------------------------------------------------- temperature

def apply_temp(dist, T):
    if T == 0:
        top = max(dist, key=dist.get)
        return {k: (1.0 if k == top else 0.0) for k in dist}
    z = {k: math.exp(math.log(v) / T) for k, v in dist.items()}
    s = sum(z.values())
    return {k: v / s for k, v in z.items()}


def entropy(dist):
    h = -sum(p * math.log2(p) for p in dist.values() if p > 0)
    return abs(h)   # avoid printing -0.0000 when one token has all the mass


def show_temp():
    print("\n\nSTEP 2  temperature, one knob\n")
    print("  Divide the logits by T before the softmax. Equivalently, raise")
    print("  each probability to the power 1/T and renormalise.\n")
    base = PEAKED
    words = list(base)
    print(f"  {'T':>6}" + "".join(f"{w:>9}" for w in words[:5])
          + f"{'entropy':>10}")
    print("  " + "-" * 62)
    for T in (0.0, 0.5, 0.8, 1.0, 1.5, 2.0):
        d = apply_temp(base, T)
        cells = "".join(f"{d[w]:>9.4f}" for w in words[:5])
        print(f"  {T:>6.1f}{cells}{entropy(d):>10.4f}")

    print(f"\n  T = 1 leaves the model's own distribution alone.")
    print(f"  T below 1 sharpens it. At T = 0 it is greedy, and the entropy")
    print(f"  is 0 because there is nothing left to choose.")
    print(f"  T above 1 flattens it and raises the entropy.")
    print(f"\n  The trade has a name at each end. Low temperature gives")
    print(f"  repetitive, safe text. High temperature gives varied text that")
    print(f"  drifts off topic and eventually stops making sense.")


# --------------------------------------------------------- top-k and top-p

def top_k(dist, k):
    keep = sorted(dist.items(), key=lambda x: -x[1])[:k]
    s = sum(v for _, v in keep)
    return {w: v / s for w, v in keep}


def top_p(dist, p):
    keep, run = [], 0.0
    for w, v in sorted(dist.items(), key=lambda x: -x[1]):
        keep.append((w, v))
        run += v
        if run >= p:
            break
    s = sum(v for _, v in keep)
    return {w: v / s for w, v in keep}


def show_truncate():
    print("\n\nSTEP 3  top-k has a problem that top-p does not\n")
    print("  Two next-token distributions. One where the model is sure, one")
    print("  where it genuinely is not.\n")
    for name, d in (("peaked", PEAKED), ("flat", FLAT)):
        top = sorted(d.items(), key=lambda x: -x[1])[:4]
        cells = ",  ".join(f"{w} {v:.2f}" for w, v in top)
        print(f"  {name:<8}{cells},  ...   entropy {entropy(d):.4f} bits")

    print(f"\n  TOP-K keeps a fixed number of candidates, whatever the shape.\n")
    print(f"  {'k':>4}{'peaked: kept':>16}{'mass kept':>12}"
          f"{'flat: kept':>14}{'mass kept':>12}")
    print("  " + "-" * 60)
    for k in (1, 2, 5, 8):
        pk = sorted(PEAKED.items(), key=lambda x: -x[1])[:k]
        fk = sorted(FLAT.items(), key=lambda x: -x[1])[:k]
        print(f"  {k:>4}{len(pk):>16}{sum(v for _, v in pk):>12.4f}"
              f"{len(fk):>14}{sum(v for _, v in fk):>12.4f}")

    print(f"\n  Read k = 5. On the peaked distribution it admits four tokens")
    print(f"  the model had all but ruled out, together worth"
          f" {sum(v for _,v in sorted(PEAKED.items(), key=lambda x:-x[1])[1:5]):.3f}.")
    print(f"  On the flat one it discards three tokens that were as good as")
    print(f"  the ones it kept.")
    print(f"\n  The same k is too loose on one shape and too tight on the")
    print(f"  other, because k cannot see the shape.")

    print(f"\n\n  TOP-P keeps the smallest set whose mass reaches p.\n")
    print(f"  {'p':>6}{'peaked: kept':>16}{'mass':>10}{'flat: kept':>14}"
          f"{'mass':>10}")
    print("  " + "-" * 58)
    for p in (0.5, 0.9, 0.95, 0.99):
        pk, fk = top_p(PEAKED, p), top_p(FLAT, p)
        mp = sum(PEAKED[w] for w in pk)
        mf = sum(FLAT[w] for w in fk)
        print(f"  {p:>6.2f}{len(pk):>16}{mp:>10.4f}{len(fk):>14}{mf:>10.4f}")

    print(f"\n  At p = 0.9 the peaked distribution keeps"
          f" {len(top_p(PEAKED, 0.9))} tokens and the flat one keeps"
          f" {len(top_p(FLAT, 0.9))}.")
    print(f"  The size of the candidate set now follows the model's own")
    print(f"  confidence, which is exactly what a fixed k could not do.")
    print(f"\n  That is why nucleus sampling replaced top-k as the default.")


# ------------------------------------------------------------ side by side

def show_compare():
    print("\n\nSTEP 4  the same distribution under every strategy\n")
    d = PEAKED
    rows = [
        ("greedy", apply_temp(d, 0.0)),
        ("temperature 0.7", apply_temp(d, 0.7)),
        ("pure sampling, T = 1", dict(d)),
        ("temperature 1.5", apply_temp(d, 1.5)),
        ("top-k, k = 3", top_k(d, 3)),
        ("top-p, p = 0.9", top_p(d, 0.9)),
    ]
    words = list(d)
    print(f"  {'strategy':<22}" + "".join(f"{w[:6]:>8}" for w in words)
          + f"{'kept':>7}{'entropy':>9}")
    print("  " + "-" * 86)
    for name, dd in rows:
        cells = "".join(f"{dd.get(w, 0.0):>8.3f}" for w in words)
        kept = sum(1 for w in words if dd.get(w, 0.0) > 0)
        print(f"  {name:<22}{cells}{kept:>7}{entropy(dd):>9.4f}")

    print(f"\n  Every row is the same model. None of them changed a weight.")
    print(f"  Decoding is a decision made after the model has spoken, and it")
    print(f"  changes the output more than most fine-tuning does.")
    print(f"\n  Which to use follows from the task. A factual answer wants")
    print(f"  the low-entropy rows. A story wants one of the middle ones.")
    print(f"  Nothing wants the bottom of the temperature range or the top.")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    for f in ("beam", "temp", "truncate", "compare"):
        ap.add_argument(f"--{f}", action="store_true")
    a = ap.parse_args()
    picked = a.beam or a.temp or a.truncate or a.compare
    if a.beam or not picked:
        show_beam()
    if a.temp or not picked:
        show_temp()
    if a.truncate or not picked:
        show_truncate()
    if a.compare or not picked:
        show_compare()
    print()


if __name__ == "__main__":
    main()