"""Generation metrics, worked as Chapter 21 works them.

Four things, all reproducing the chapter's tables:

    modified precision  why plain n-gram precision is gameable
    the brevity penalty why BLEU needs a length term, and what it costs
    BLEU end to end     all four n, the geometric mean, the final score
    ROUGE and the gap   recall instead of precision, and what neither sees

    python3 bleu.py               # all four
    python3 bleu.py --clip        # the clipping counter-example
    python3 bleu.py --brevity     # the length term
    python3 bleu.py --bleu        # a full score, step by step
    python3 bleu.py --rouge       # ROUGE, and where both metrics fail

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

import argparse
import math
from collections import Counter


def ngrams(tokens, n):
    return Counter(tuple(tokens[i:i + n]) for i in range(len(tokens) - n + 1))


def modified_precision(cand, refs, n):
    """Clipped precision: no n-gram may be credited more often than the
    best single reference contains it."""
    c = ngrams(cand, n)
    if not c:
        return 0, 0
    ceiling = Counter()
    for r in refs:
        rc = ngrams(r, n)
        for g in rc:
            ceiling[g] = max(ceiling[g], rc[g])
    kept = sum(min(count, ceiling[g]) for g, count in c.items())
    return kept, sum(c.values())


def brevity_penalty(cand, refs):
    c = len(cand)
    r = min((len(x) for x in refs), key=lambda L: (abs(L - c), L))
    return (1.0 if c > r else math.exp(1 - r / c) if c else 0.0), c, r


def bleu(cand, refs, N=4):
    ps = []
    for n in range(1, N + 1):
        kept, total = modified_precision(cand, refs, n)
        ps.append((kept, total, kept / total if total else 0.0))
    bp, c, r = brevity_penalty(cand, refs)
    if any(p == 0 for _, _, p in ps):
        return 0.0, ps, bp, c, r
    logmean = sum(math.log(p) for _, _, p in ps) / N
    return bp * math.exp(logmean), ps, bp, c, r


# ---------------------------------------------------------------- clipping

def show_clip():
    print("\nSTEP 1  why precision has to be clipped\n")
    refs = ["the cat is on the mat".split(),
            "there is a cat on the mat".split()]
    print("  references:")
    for r in refs:
        print(f"    '{' '.join(r)}'")

    cheat = "the the the the the the the".split()
    honest = "the cat is on the mat".split()

    print(f"\n  A degenerate candidate: '{' '.join(cheat)}'\n")
    c = ngrams(cheat, 1)
    print(f"  It contains 'the' {c[('the',)]} times, and every one of them")
    print(f"  does appear in a reference. Unclipped unigram precision is")
    print(f"  {c[('the',)]}/{sum(c.values())} = 1.0000, a perfect score for")
    print(f"  a sentence that says nothing.\n")

    kept, total = modified_precision(cheat, refs, 1)
    print(f"  Clipping caps each n-gram at the most any single reference")
    print(f"  contains. Reference 1 has 'the' twice, so the ceiling is 2.\n")
    print(f"  {'candidate':<32}{'clipped':>9}{'total':>8}{'p_1':>9}")
    print("  " + "-" * 60)
    for name, cand in (("the the the ... (7 times)", cheat),
                       ("the cat is on the mat", honest)):
        k, t = modified_precision(cand, refs, 1)
        print(f"  {name:<32}{k:>9}{t:>8}{k / t:>9.4f}")

    print(f"\n  The degenerate candidate drops to {kept}/{total} ="
          f" {kept / total:.4f}.")
    print(f"  Clipping is not a refinement. Without it the metric can be won")
    print(f"  by repeating one word, so it would measure nothing at all.")


# ---------------------------------------------------------- brevity penalty

def show_brevity():
    print("\n\nSTEP 2  precision alone rewards saying less\n")
    refs = ["the cat is on the mat".split()]
    print(f"  reference: '{' '.join(refs[0])}'   length {len(refs[0])}\n")
    print(f"  {'candidate':<26}{'len':>5}{'p_1':>9}{'BP':>9}"
          f"{'p_1 x BP':>11}")
    print("  " + "-" * 60)
    for cand in ("the", "the cat", "the cat is on", "the cat is on the mat",
                 "the cat is on the mat today for a while"):
        toks = cand.split()
        k, t = modified_precision(toks, refs, 1)
        bp, c, r = brevity_penalty(toks, refs)
        p = k / t if t else 0.0
        print(f"  {cand:<26}{len(toks):>5}{p:>9.4f}{bp:>9.4f}{p * bp:>11.4f}")

    print(f"\n  Look at the first row. One word, perfectly precise, and")
    print(f"  useless. Precision alone would rank it top.")
    print(f"\n  The brevity penalty is exp(1 - r/c) when the candidate is")
    print(f"  shorter than the reference, and 1 otherwise:\n")
    print(f"    BP = 1               if c > r")
    print(f"    BP = exp(1 - r/c)    if c <= r\n")
    print(f"  It is deliberately asymmetric. Too short is punished, too long")
    print(f"  is not, because a long candidate is already punished by")
    print(f"  precision: the extra tokens are unmatched and dilute the ratio.")
    print(f"\n  Read the last row. Nine tokens against six, BP stays at 1.00,")
    print(f"  and the score falls anyway because precision fell.")


# --------------------------------------------------------------- full BLEU

def show_bleu():
    refs = ["the cat is on the mat".split(),
            "there is a cat on the mat".split()]
    cands = {
        "the cat is on the mat": "the cat is on the mat",
        "a cat is on the mat": "a cat is on the mat",
        "the cat sat on the mat": "the cat sat on the mat",
        "on mat the cat is the": "on mat the cat is the",
    }

    print("\n\nSTEP 3  BLEU end to end\n")
    print("  references:")
    for r in refs:
        print(f"    '{' '.join(r)}'")
    print()

    for label, text in cands.items():
        cand = text.split()
        score, ps, bp, c, r = bleu(cand, refs)
        print(f"  candidate '{label}'")
        print(f"    {'n':>3}{'clipped':>9}{'total':>7}{'p_n':>9}")
        print("    " + "-" * 30)
        for n, (kept, total, p) in enumerate(ps, start=1):
            print(f"    {n:>3}{kept:>9}{total:>7}{p:>9.4f}")
        print(f"    candidate length {c}, closest reference {r}, BP {bp:.4f}")
        print(f"    BLEU-4 = {score:.4f}\n")

    print("  Three things that table shows.\n")
    print("  BLEU uses the GEOMETRIC mean of the four precisions. One zero")
    print("  anywhere sends the whole score to zero, which is why the last")
    print("  candidate scores 0 despite getting most unigrams right.")
    print("\n  That is severe on short texts and it is the reason BLEU is")
    print("  reported over a corpus rather than a sentence.")
    print("\n  The scrambled candidate contains exactly the right words in")
    print("  the wrong order. Its unigram precision is high and its higher")
    print("  n-grams collapse, which is the only way BLEU sees word order.")


# ------------------------------------------------------------------- ROUGE

def rouge_n(cand, refs, n):
    c = ngrams(cand, n)
    best = 0.0
    for r in refs:
        rc = ngrams(r, n)
        if not rc:
            continue
        overlap = sum(min(c[g], rc[g]) for g in rc)
        best = max(best, overlap / sum(rc.values()))
    return best


def show_rouge():
    refs = ["the cat is on the mat".split()]
    print("\n\nSTEP 4  ROUGE measures the other direction\n")
    print(f"  reference: '{' '.join(refs[0])}'\n")
    print("  BLEU asks: how much of what I said was in the reference?")
    print("  ROUGE asks: how much of the reference did I say?\n")
    print("  Precision against recall, and summarisation cares about the")
    print("  second one, because leaving things out is the failure mode.\n")

    print(f"  {'candidate':<44}{'BLEU-4':>9}{'ROUGE-1':>10}{'ROUGE-2':>10}")
    print("  " + "-" * 74)
    for text in ("the cat is on the mat",
                 "the cat",
                 "the cat is on the mat and also on a fine rug",
                 "the feline rests upon the rug"):
        cand = text.split()
        score, *_ = bleu(cand, refs)
        print(f"  {text:<44}{score:>9.4f}{rouge_n(cand, refs, 1):>10.4f}"
              f"{rouge_n(cand, refs, 2):>10.4f}")

    print(f"\n  Row 2 is short. BLEU tolerates it less than ROUGE does,")
    print(f"  because ROUGE-1 still counts the two words it did recover.")
    print(f"\n  Row 3 is padded. ROUGE gives it a perfect 1.0000, because")
    print(f"  every reference n-gram is present. BLEU does not, because the")
    print(f"  padding is unmatched. Recall alone can be gamed by saying")
    print(f"  everything, which is why ROUGE is usually reported with its")
    print(f"  F-measure variant.")
    print(f"\n  Row 4 is the one that matters. It is a correct paraphrase.")
    print(f"  BLEU gives it 0.0000 and ROUGE-2 gives it 0.0000. ROUGE-1")
    print(f"  gives it 0.3333, and every point of that comes from the word")
    print(f"  'the'. Not one content word was credited.")
    print(f"\n  Neither metric has any notion of meaning. They count string")
    print(f"  overlap, and a paraphrase overlaps in the wrong places.")
    print(f"\n  That single row is the argument for everything that came")
    print(f"  after: embedding-based metrics, learned metrics, and human")
    print(f"  evaluation.")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    for f in ("clip", "brevity", "bleu", "rouge"):
        ap.add_argument(f"--{f}", action="store_true")
    a = ap.parse_args()
    picked = a.clip or a.brevity or a.bleu or a.rouge
    if a.clip or not picked:
        show_clip()
    if a.brevity or not picked:
        show_brevity()
    if a.bleu or not picked:
        show_bleu()
    if a.rouge or not picked:
        show_rouge()
    print()


if __name__ == "__main__":
    main()