"""The recurrent network, worked as Chapter 12 works it.

Four things, all reproducing the chapter's tables:

    unrolling          three time steps by hand, h_0 to h_3
    the parameter count  why an RNN does not grow with the sequence
    vanishing gradients  the product of Jacobians, and where it dies
    gradient clipping    what the rescale actually does

    python3 rnn.py               # all four
    python3 rnn.py --unroll      # the three-step trace
    python3 rnn.py --params      # RNN against a feed-forward window
    python3 rnn.py --vanish      # the decay table, and the 47-word example
    python3 rnn.py --clip        # clipping, before and after

Install:  pip install numpy
"""

import argparse
import math

try:
    import numpy as np
except ImportError:
    raise SystemExit("this one needs numpy:  pip install numpy")


# ------------------------------------------------------------- the recurrence

# A model small enough to print in full. Two hidden units, three vocabulary
# words, embeddings of size two.
WORDS = ["the", "dog", "ran"]
E = np.array([[0.50, -0.20],       # the
              [0.10, 0.80],        # dog
              [-0.60, 0.30]])      # ran
W = np.array([[0.60, -0.30],       # input to hidden, 2 x 2
              [0.20, 0.50]])
U = np.array([[0.40, 0.10],        # hidden to hidden, 2 x 2
              [-0.20, 0.70]])
bh = np.array([0.05, -0.05])
V = np.array([[0.70, -0.30],       # hidden to vocabulary, 3 x 2
              [-0.40, 0.60],
              [0.10, 0.20]])
by = np.zeros(3)


def softmax(z):
    e = np.exp(z - z.max())
    return e / e.sum()


def show_unroll():
    print("\nSTEP 1  unrolling three time steps by hand\n")
    print("  h_t = tanh(W x_t + U h_{t-1} + b_h),   y_t = softmax(V h_t + b_y)")
    print(f"\n  The same W, U and V are used at every step. That is what the")
    print(f"  word 'recurrent' means, and it is why the parameter count does")
    print(f"  not grow with the sentence.\n")

    seq = ["the", "dog", "ran"]
    h = np.zeros(2)
    print(f"  h_0 = {h}   the state before any word is read\n")
    def vec(a):
        return "(" + ", ".join(f"{v:+.4f}" for v in a) + ")"

    print(f"  {'t':>2}{'word':>6}{'x_t':>21}{'W x + U h + b':>21}{'h_t':>21}")
    print("  " + "-" * 71)
    for t, w in enumerate(seq, start=1):
        x = E[WORDS.index(w)]
        pre = W @ x + U @ h + bh
        h = np.tanh(pre)
        print(f"  {t:>2}{w:>6}{vec(x):>21}{vec(pre):>21}{vec(h):>21}")

    print(f"\n  Every h_t carries the whole prefix, not a fixed window.")
    print(f"  h_3 was computed from h_2, which was computed from h_1.")

    print(f"\n  The output at each step is a distribution over the vocabulary:\n")
    h = np.zeros(2)
    print(f"  {'t':>2}{'read':>6}   " + "".join(f"{w:>9}" for w in WORDS)
          + f"{'  most likely next':>20}")
    print("  " + "-" * 62)
    for t, w in enumerate(seq, start=1):
        h = np.tanh(W @ E[WORDS.index(w)] + U @ h + bh)
        p = softmax(V @ h + by)
        print(f"  {t:>2}{w:>6}   " + "".join(f"{v:>9.4f}" for v in p)
              + f"{WORDS[int(p.argmax())]:>20}")
    print("\n  The weights are random here, so the predictions mean nothing.")
    print("  The shape is the point. One distribution per time step, and a")
    print("  language model needs exactly one distribution per time step.")


# ------------------------------------------------------------ the parameters

def show_params(V_size=10000, d=100, H=500):
    print(f"\n\nSTEP 2  the parameter count stops depending on the sequence\n")
    print(f"  |V| = {V_size},  embedding d = {d},  hidden H = {H}\n")
    print(f"  {'matrix':<10}{'shape':<18}{'parameters':>14}")
    print("  " + "-" * 44)
    rows = [("E", f"{V_size} x {d}", V_size * d),
            ("W", f"{H} x {d}", H * d),
            ("U", f"{H} x {H}", H * H),
            ("V", f"{H} x {V_size}", H * V_size),
            ("biases", f"{H} + {V_size}", H + V_size)]
    total = sum(n for _, _, n in rows)
    for name, shape, n in rows:
        print(f"  {name:<10}{shape:<18}{n:>14,}")
    print("  " + "-" * 44)
    print(f"  {'total':<28}{total:>14,}")

    print(f"\n  Not one of those shapes mentions the sentence length.")
    print(f"  U is H x H whether the sentence has 5 words or 500.\n")

    print(f"  Compare the feed-forward model of Chapter 11, whose W_h is")
    print(f"  (context x d) x H and therefore grows with the window:\n")
    print(f"  {'context words':>15}{'feed-forward W_h':>20}{'RNN U':>12}")
    print("  " + "-" * 48)
    for ctx in (3, 5, 10, 50, 500):
        print(f"  {ctx:>15}{ctx * d * H:>20,}{H * H:>12,}")

    print(f"\n  The left column grows without limit. The right column is a")
    print(f"  constant. That single fact is why the recurrence was worth")
    print(f"  inventing, and it is what lets an RNN read a whole paragraph.")


# ------------------------------------------------------- vanishing gradients

def show_vanish():
    print("\n\nSTEP 3  why the memory fades\n")
    print("  Backpropagation through time multiplies one Jacobian per step:\n")
    print("    dE/dh_0  =  dE/dh_tau  *  product over t of  dh_t / dh_{t-1}\n")
    print("  Each factor is diag(tanh'(...)) U. So the whole product behaves")
    print("  like a number raised to the power of the distance.\n")

    print(f"  {'distance':>9}" + "".join(f"{f'x {g}':>13}" for g in
                                         (0.5, 0.9, 1.0, 1.1, 1.2)))
    print("  " + "-" * 74)
    for tau in (1, 5, 10, 20, 47, 100):
        cells = "".join(f"{g ** tau:>13.3e}" for g in (0.5, 0.9, 1.0, 1.1, 1.2))
        print(f"  {tau:>9}{cells}")

    print(f"\n  Read the columns. Below 1 the gradient dies. Above 1 it")
    print(f"  explodes. Exactly at 1 it survives, and nothing keeps it there.")
    print(f"\n  The threshold is the largest eigenvalue of U. If every")
    print(f"  eigenvalue satisfies |lambda| < 1 the gradients vanish, and if")
    print(f"  any satisfies |lambda| > 1 they explode.")

    print(f"\n\n  THE 47 WORD EXAMPLE\n")
    print("  'Raj entered CoffeeDay to meet his partner Dru. ... After a long")
    print("  and fruitful discussion, Raj said goodbye to his ______'")
    print("\n  The answer is 'partner', 47 words back. What reaches it:\n")
    for g in (0.01, 0.5, 1.2):
        v = g ** 47
        verdict = ("gone entirely" if v < 1e-12
                   else "exploded" if v > 1 else "too small to matter")
        print(f"    {g} ^ 47 = {v:.6e}   {verdict}")

    print("\n  A gradient of 1e-94 is not a small update. It is no update.")
    print("  The network cannot learn the dependency, however true it is.")
    print("\n  Note what is not broken. The forward pass carries information")
    print("  fine. It is the backward pass that cannot deliver the blame,")
    print("  so the weights that would capture the link never move.")


# ---------------------------------------------------------- gradient clipping

def show_clip(threshold=5.0):
    print(f"\n\nSTEP 4  gradient clipping   (threshold = {threshold})\n")
    print("  Vanishing needs a new architecture. Exploding has a cheap fix.")
    print("  If the gradient is longer than the threshold, rescale it to the")
    print("  threshold and keep its direction:\n")
    print("    if ||g|| > threshold:   g <- (threshold / ||g||) g\n")
    rng = np.random.default_rng(0)
    print(f"  {'step':>5}{'||g|| before':>15}{'||g|| after':>14}"
          f"{'direction kept':>17}")
    print("  " + "-" * 52)
    for step, scale in enumerate([0.4, 1.2, 3.0, 40.0, 900.0], start=1):
        g = rng.normal(0, 1, 6) * scale
        n = np.linalg.norm(g)
        g2 = g * (threshold / n) if n > threshold else g
        cos = float(g @ g2 / (np.linalg.norm(g) * np.linalg.norm(g2)))
        print(f"  {step:>5}{n:>15.4f}{np.linalg.norm(g2):>14.4f}"
              f"{cos:>17.4f}")

    print("\n  The last column is 1.0000 every time. Clipping changes how far")
    print("  the step goes and never which way it points.")
    print("\n  That is the whole trick, and it is why one line of code stops")
    print("  a training run from diverging. It does nothing at all for the")
    print("  vanishing case, where the gradient is already too small.")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    for f in ("unroll", "params", "vanish", "clip"):
        ap.add_argument(f"--{f}", action="store_true")
    a = ap.parse_args()
    picked = a.unroll or a.params or a.vanish or a.clip
    if a.unroll or not picked:
        show_unroll()
    if a.params or not picked:
        show_params()
    if a.vanish or not picked:
        show_vanish()
    if a.clip or not picked:
        show_clip()
    print()


if __name__ == "__main__":
    main()