"""n-gram language models, worked as Chapter 10 works them.

Every table in the chapter comes from this file. The conventions match
`autograder/ngram_reference.py` exactly, so the numbers here are the numbers
Assignment 4 grades against:

    pad with (n-1) copies of <s> and one </s>
    V holds every training word, plus </s> and <unk>
    <s> is context only, never predicted, never mapped to <unk>
    perplexity is over the predicted tokens: the words plus </s>

    python3 ngram_lm.py                  # every step
    python3 ngram_lm.py --counts         # counts, MLE, add-one, side by side
    python3 ngram_lm.py --perplexity     # the four-row perplexity table
    python3 ngram_lm.py --alpha          # the U shape that picks alpha
    python3 ngram_lm.py --params         # why n = 5 is not an option
    python3 ngram_lm.py --interpolate    # backoff and interpolation
    python3 ngram_lm.py --unk            # what happens to an unseen word

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

import argparse
import math
from collections import Counter

BOS, EOS, UNK = "<s>", "</s>", "<unk>"

TRAIN = ["the cat sat", "the cat ran", "the dog sat"]

# A second corpus for interpolation, where the trigram counts run out but the
# bigram counts do not.
TRAIN3 = [
    "the cat sat on the mat",
    "the cat sat on the rug",
    "the dog sat on the mat",
    "a cat ran to the mat",
]


def tokenize(text):
    return text.lower().split()


def padded(tokens, n):
    return [BOS] * (n - 1) + list(tokens) + [EOS]


class NGramLM:
    """Add-alpha smoothed n-gram model. Same arithmetic as the autograder."""

    def __init__(self, n=2, alpha=1.0):
        self.n, self.alpha = n, alpha
        self.vocab = set()
        self.ngram = Counter()
        self.context = Counter()

    def train(self, sentences):
        for s in sentences:
            self.vocab.update(tokenize(s))
        self.vocab.update([EOS, UNK])
        for s in sentences:
            p = padded(tokenize(s), self.n)
            for i in range(self.n - 1, len(p)):
                ctx = tuple(p[i - self.n + 1:i])
                self.ngram[ctx + (p[i],)] += 1
                self.context[ctx] += 1
        return self

    def _map(self, w):
        # <s> is padding, not a word. It is deliberately outside V, so sending
        # it to <unk> would throw away every start-of-sentence statistic.
        return w if w == BOS or w in self.vocab else UNK

    def prob(self, word, context):
        ctx = tuple(self._map(c) for c in context)
        num = self.ngram.get(ctx + (self._map(word),), 0) + self.alpha
        den = self.context.get(ctx, 0) + self.alpha * len(self.vocab)
        return num / den if den else 0.0

    def scored(self, sentence):
        """Every (context, word, probability) the sentence is scored on."""
        p = padded(tokenize(sentence), self.n)
        out = []
        for i in range(self.n - 1, len(p)):
            ctx = tuple(p[i - self.n + 1:i])
            out.append((ctx, p[i], self.prob(p[i], ctx)))
        return out

    def perplexity(self, sentence):
        rows = self.scored(sentence)
        lp = sum(math.log2(p) for _, _, p in rows)
        return 2 ** (-lp / len(rows))

    def corpus_perplexity(self, sentences):
        lp = n = 0.0
        for s in sentences:
            rows = self.scored(s)
            lp += sum(math.log2(p) for _, _, p in rows)
            n += len(rows)
        return 2 ** (-lp / n)


def mle(sentences, word, given, n=2):
    """Unsmoothed maximum likelihood estimate, straight from the counts."""
    ng, ctx = Counter(), Counter()
    for s in sentences:
        p = padded(tokenize(s), n)
        for i in range(n - 1, len(p)):
            c = tuple(p[i - n + 1:i])
            ng[c + (p[i],)] += 1
            ctx[c] += 1
    g = given if isinstance(given, tuple) else (given,)
    return ng[g + (word,)] / ctx[g] if ctx[g] else 0.0


def raw_counts(sentences, n=2):
    ng, ctx = Counter(), Counter()
    for s in sentences:
        p = padded(tokenize(s), n)
        for i in range(n - 1, len(p)):
            c = tuple(p[i - n + 1:i])
            ng[c + (p[i],)] += 1
            ctx[c] += 1
    return ng, ctx


# ------------------------------------------------------------------ step one

def show_counts():
    m = NGramLM(2, 1.0).train(TRAIN)
    ng, ctx = raw_counts(TRAIN)
    V = sorted(m.vocab)

    print("\nSTEP 1  three sentences, and what they say about 'the'\n")
    print("  training:  " + "   ".join(f"'{s}'" for s in TRAIN))
    print(f"\n  vocabulary V = {V}")
    print(f"  |V| = {len(m.vocab)}.  <s> is not in V. It is context, never a"
          " prediction.")

    print("\n  contexts and how often each was seen:")
    print("    " + "   ".join(f"{c[0]}: {n}" for c, n in sorted(ctx.items())))

    for given in ("the", "dog"):
        print(f"\n\n  CONTEXT '{given}'   seen {ctx[(given,)]} times\n")
        print(f"  {'next word':<12}{'count':>7}{'MLE':>10}{'add-1':>10}"
              f"{'  what add-1 did'}")
        print("  " + "-" * 62)
        tm = ta = 0.0
        for w in V:
            c = ng[(given, w)]
            a, b = mle(TRAIN, w, given), m.prob(w, (given,))
            tm, ta = tm + a, ta + b
            note = "gave it mass" if a == 0 else f"took {a - b:+.4f}"
            print(f"  {w:<12}{c:>7}{a:>10.4f}{b:>10.4f}   {note}")
        print("  " + "-" * 62)
        print(f"  {'total':<12}{ctx[(given,)]:>7}{tm:>10.4f}{ta:>10.4f}")

    print("\n  Both columns sum to 1, which is the point of the alpha|V| term")
    print("  in the denominator. The counts did not change. The belief did.")
    print("\n  Read the 'dog' table again. MLE says p(sat|dog) = 1, so the")
    print("  model believes 'the dog ran' is impossible. Add-1 drops that")
    print("  certainty to 0.25 and hands 0.125 to every other word.")


# ------------------------------------------------------------------ step two

def show_perplexity(sentence="the cat sat", alpha=1.0):
    m = NGramLM(2, alpha).train(TRAIN)
    ng, ctx = raw_counts(TRAIN)
    rows = m.scored(sentence)

    print(f"\n\nSTEP 2  perplexity of '{sentence}'   (bigram, alpha = {alpha})\n")
    print(f"  {'predicted':<12}{'context':<10}{'count':>7}{'of':>5}"
          f"{'p add-1':>10}{'log2 p':>10}")
    print("  " + "-" * 56)
    total = 0.0
    for c, w, p in rows:
        total += math.log2(p)
        print(f"  {w:<12}{c[0]:<10}{ng[c + (w,)]:>7}{ctx[c]:>5}"
              f"{p:>10.4f}{math.log2(p):>10.4f}")
    n = len(rows)
    h = -total / n
    print("  " + "-" * 56)
    print(f"  {'sum':<34}{'':>10}{total:>10.4f}")
    print(f"\n  N = {n} predicted tokens. The words, plus one </s>.")
    print(f"  <s> was context for the first row and was never predicted.")
    print(f"\n  H = -({total:.4f}) / {n} = {h:.4f} bits per token")
    print(f"  PP = 2^{h:.4f} = {2 ** h:.4f}")
    print(f"\n  So the model is about as unsure at each step as if it were")
    print(f"  choosing uniformly among {2 ** h:.2f} words. The vocabulary has"
          f" {len(m.vocab)}.")


# ---------------------------------------------------------------- step three

def show_alpha():
    seen, held = "the cat sat", "the dog ran"
    grid = [0.01, 0.02, 0.05, 0.1, 0.15, 0.2, 0.3, 0.5, 1.0, 2.0]

    print("\n\nSTEP 3  choosing alpha, and why the test set has to be new\n")
    print(f"  '{seen}' appears in training. '{held}' does not.")
    print(f"  Every word of '{held}' does. Only the bigram (dog, ran) is new.\n")
    print(f"  {'alpha':>7}{'PP seen':>12}{'PP held out':>14}"
          f"{'p(cat|the)':>13}")
    print("  " + "-" * 46)
    best = (float("inf"), None)
    for a in grid:
        m = NGramLM(2, a).train(TRAIN)
        p_seen, p_held = m.perplexity(seen), m.perplexity(held)
        if p_held < best[0]:
            best = (p_held, a)
        print(f"  {a:>7}{p_seen:>12.4f}{p_held:>14.4f}"
              f"{m.prob('cat', ('the',)):>13.4f}")

    print(f"\n  The seen column falls all the way down. Less smoothing always")
    print(f"  looks better on text the model memorised.")
    print(f"\n  The held out column is a U. It bottoms at alpha = {best[1]},")
    print(f"  PP = {best[0]:.4f}, and rises on both sides.")
    print(f"\n  Too little smoothing and the one unseen bigram is crushed.")
    print(f"  Too much and the seen bigrams are robbed to pay for it.")
    worst = NGramLM(2, grid[0]).train(TRAIN).perplexity(held)
    print(f"\n  Tuning alpha on training text would have picked {grid[0]}, the")
    print(f"  bottom of the seen column. On held out text that scores"
          f" {worst:.4f}")
    print(f"  against the best available {best[0]:.4f}, so it is"
          f" {worst / best[0]:.2f} times worse.")


# ----------------------------------------------------------------- step four

def show_params():
    print("\n\nSTEP 4  the curse of dimensionality, in one table\n")
    print("  An n-gram model needs one number per (history, word) pair.")
    print("  There are |V|^(n-1) histories and |V|-1 free choices in each.\n")
    print(f"  {'|V|':>8}{'n=1':>12}{'n=2':>12}{'n=3':>12}{'n=4':>12}"
          f"{'n=5':>12}")
    print("  " + "-" * 68)
    for V in (1000, 10000, 50000):
        cells = "".join(f"{V ** (n - 1) * (V - 1):>12.1e}" for n in range(1, 6))
        print(f"  {V:>8}" + cells)

    V, n = 50000, 5
    params = V ** (n - 1) * (V - 1)
    corpus = 1e12                      # a trillion tokens, a large corpus
    print(f"\n  At |V| = {V} a 5-gram model has {params:.3e} parameters.")
    print(f"\n  Now count what could ever fill them. A corpus of {corpus:.0e}")
    print(f"  tokens contains at most {corpus:.0e} distinct 5-grams, one per")
    print(f"  position. So at most {corpus / params:.1e} of the cells can hold")
    print(f"  a nonzero count. That is about one cell in"
          f" {params / corpus:.0e}.")
    print("\n  The counts do not merely get thin. Almost every cell is empty,")
    print("  and no corpus that will ever exist can fill them.")
    print("\n  This is the wall Chapter 11 walks into, and the reason neural")
    print("  language models exist. They do not store a cell per history.")


# ----------------------------------------------------------------- step five

def show_interpolate():
    ng3, ctx3 = raw_counts(TRAIN3, 3)
    ng2, ctx2 = raw_counts(TRAIN3, 2)
    ng1, ctx1 = raw_counts(TRAIN3, 1)
    lam = (0.1, 0.3, 0.6)      # unigram, bigram, trigram

    print("\n\nSTEP 5  backoff and interpolation\n")
    print("  training:")
    for s in TRAIN3:
        print(f"    '{s}'")
    print(f"\n  weights: lambda1 = {lam[0]} unigram, lambda2 = {lam[1]}"
          f" bigram, lambda3 = {lam[2]} trigram")
    print(f"  they sum to {sum(lam)}, which is the one constraint.\n")

    # Four queries chosen to walk down the staircase: all orders agree, then
    # the trigram sharpens, then the trigram fails, then the bigram fails too.
    queries = [("on", ("cat", "sat")), ("mat", ("on", "the")),
               ("cat", ("on", "the")), ("ran", ("on", "the"))]
    print(f"  {'query':<22}{'p3 tri':>9}{'p2 bi':>9}{'p1 uni':>9}"
          f"{'mixed':>10}   seen as")
    print("  " + "-" * 74)
    for w, ctx in queries:
        p3 = mle(TRAIN3, w, ctx, 3)
        p2 = mle(TRAIN3, w, ctx[-1], 2)
        p1 = ng1[(w,)] / sum(ctx1.values())
        mix = lam[0] * p1 + lam[1] * p2 + lam[2] * p3
        note = (f"trigram {ng3[ctx + (w,)]}x, bigram {ng2[(ctx[-1], w)]}x")
        print(f"  {'p(' + w + ' | ' + ' '.join(ctx) + ')':<22}"
              f"{p3:>9.4f}{p2:>9.4f}{p1:>9.4f}{mix:>10.4f}   {note}")

    print("\n  Row 1, 'cat sat on'. Every order is certain and the mix is high.")
    print("  Row 2, 'on the mat'. The trigram is sharper than the bigram,")
    print("  because it has the extra word of context, and the mix lands")
    print("  between them.")
    print("\n  Row 3 is the one that matters. 'on the cat' never occurred, so")
    print("  the trigram says impossible. The bigram has seen 'the cat' twice")
    print("  and says 0.2857. The mix keeps the sentence alive.")
    print("\n  Row 4, 'ran' never follows 'the' at all. Both higher orders")
    print("  fail and only the unigram is left. The answer is small but not")
    print("  zero, which is the entire purpose.")
    print("\n  Backoff does the same job with a switch instead of a blend.")
    print("  Use the trigram if you have seen it, otherwise drop an order.")


# ------------------------------------------------------------------ step six

def show_unk():
    m = NGramLM(2, 1.0).train(TRAIN)
    print("\n\nSTEP 6  a word the model has never met\n")
    tests = ["the cat sat", "the dog ran", "the cat slept", "a zebra danced"]
    print(f"  {'sentence':<18}{'maps to':<34}{'PP':>9}")
    print("  " + "-" * 62)
    for s in tests:
        mapped = " ".join(m._map(w) for w in tokenize(s))
        print(f"  {s:<18}{mapped:<34}{m.perplexity(s):>9.4f}")
    print("\n  Unseen words become <unk>, which is an ordinary member of V.")
    print("  So the model has a probability for them and never returns zero.")
    print("  Perplexity still rises, and it should. The model really is more")
    print("  surprised by a sentence it cannot read.")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    for flag in ("counts", "perplexity", "alpha", "params", "interpolate",
                 "unk"):
        ap.add_argument(f"--{flag}", action="store_true")
    ap.add_argument("--sentence", default="the cat sat")
    ap.add_argument("--alpha-value", type=float, default=1.0)
    a = ap.parse_args()

    picked = any([a.counts, a.perplexity, a.alpha, a.params, a.interpolate,
                  a.unk])
    if a.counts or not picked:
        show_counts()
    if a.perplexity or not picked:
        show_perplexity(a.sentence, a.alpha_value)
    if a.alpha or not picked:
        show_alpha()
    if a.params or not picked:
        show_params()
    if a.interpolate or not picked:
        show_interpolate()
    if a.unk or not picked:
        show_unk()
    print()


if __name__ == "__main__":
    main()