"""HAL, the Hyperspace Analogue to Language, traced as Chapter 7 traces it.

Two scans over the corpus, a ramped window, and an asymmetric matrix. Run this
and you get the same matrix the book prints for

    the horse raced past the barn fell

    python3 hal.py                      # the book's corpus, window 5
    python3 hal.py --window 3           # watch the ramp shorten
    python3 hal.py --corpus "a b a c"   # your own text
    python3 hal.py --vectors            # concatenated row+column vectors

The one idea to hold on to: M[a][b] counts b occurring BEFORE a, so M[b][a]
counts something different. The matrix is asymmetric on purpose, because word
order carries grammar.

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

import argparse
from collections import Counter

CORPUS = "the horse raced past the barn fell"


def ramp(window, distance):
    """Weight for a neighbour `distance` words away. Equation 7.1."""
    return window - distance + 1


def build(tokens, window):
    """M[a][b] = summed weight of b occurring BEFORE a, within the window.

    This single matrix holds both scans. Reading across a row gives you the
    left-to-right information (what preceded this word). Reading down a column
    gives you the right-to-left information (what followed it).
    """
    vocab = list(dict.fromkeys(tokens))   # first-appearance order, as the book
    M = {a: {b: 0 for b in vocab} for a in vocab}
    for i, target in enumerate(tokens):
        for j in range(max(0, i - window), i):
            M[target][tokens[j]] += ramp(window, i - j)
    return vocab, M


def show_scan(tokens, window):
    """The ramp as the lectures draw it, one key word per row."""
    print(f"\nTHE RAMP  (window {window}: nearest neighbour scores {window},"
          f" furthest scores 1)\n")
    print("       " + "".join(f"{t:>7}" for t in tokens))
    for k in range(min(2, len(tokens))):
        cells = []
        for i in range(len(tokens)):
            if i < k:
                cells.append("")
            elif i == k:
                cells.append("K")
            else:
                d = i - k
                cells.append(str(ramp(window, d)) if d <= window else "0")
        print(f"  from {tokens[k]:<3}" + "".join(f"{c:>7}" for c in cells))
    print("\n  K marks the key word. Each following word gets a smaller weight.")


def show_matrix(vocab, M, tokens):
    print(f"\nTHE MATRIX  M[row][col] = weight of COLUMN word before ROW word\n")
    print("            " + "".join(f"{b:>11}" for b in vocab))
    for a in vocab:
        print(f"  {a:<10}" + "".join(f"{M[a][b]:>11}" for b in vocab))

    # The cell where a repeated context word accumulates twice.
    repeats = [w for w, n in Counter(tokens).items() if n > 1]
    if repeats:
        r = repeats[0]
        best = max(vocab, key=lambda a: M[a][r])
        if M[best][r]:
            print(f"\n  '{r}' occurs {Counter(tokens)[r]} times, so it can")
            print(f"  contribute more than once. Row '{best}' collects"
                  f" {M[best][r]} from it.")


def show_asymmetry(vocab, M):
    print("\nWHY IT IS ASYMMETRIC")
    print("  M[a][b] counts b BEFORE a.  M[b][a] counts a BEFORE b.")
    print("  Different events, so the two cells disagree.\n")
    pairs = []
    order = sorted(vocab)
    for i, a in enumerate(order):
        for b in order[i + 1:]:
            if M[a][b] != M[b][a]:
                pairs.append((abs(M[a][b] - M[b][a]), a, b))
    pairs.sort(reverse=True)
    print(f"  {'pair':<22}{'M[a][b]':>9}{'M[b][a]':>9}")
    for _, a, b in pairs[:5]:
        print(f"  {a + ', ' + b:<22}{M[a][b]:>9}{M[b][a]:>9}")
    if not pairs:
        print("  (none: this corpus happens to be symmetric)")


def vectors(vocab, M):
    """A word's full HAL vector: its row, then its column. Length 2|V|."""
    return {a: [M[a][b] for b in vocab] + [M[b][a] for b in vocab]
            for a in vocab}


def minkowski(x, y, r=2):
    """Equation 7.5. r=2 is ordinary Euclidean distance."""
    return sum(abs(p - q) ** r for p, q in zip(x, y)) ** (1 / r)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--corpus", default=CORPUS)
    ap.add_argument("--window", type=int, default=5)
    ap.add_argument("--vectors", action="store_true",
                    help="print the concatenated row+column vectors")
    a = ap.parse_args()

    tokens = a.corpus.split()
    print(f"\nCORPUS  {a.corpus}")
    vocab, M = build(tokens, a.window)
    show_scan(tokens, a.window)
    show_matrix(vocab, M, tokens)
    show_asymmetry(vocab, M)

    if a.vectors:
        V = vectors(vocab, M)
        print(f"\nFULL VECTORS  row then column, {2 * len(vocab)} dimensions")
        for w in vocab:
            row = " ".join(f"{x:>2}" for x in V[w][:len(vocab)])
            col = " ".join(f"{x:>2}" for x in V[w][len(vocab):])
            print(f"  {w:<10}[{row}] + [{col}]")
        print("\n  The row is what preceded the word. The column is what")
        print("  followed it. A symmetric model would have nothing to join.")

        print(f"\n  Minkowski distance (r=2) between a few pairs:")
        for x, y in [(vocab[0], vocab[1]), (vocab[0], vocab[-1])]:
            print(f"    {x:<8} to {y:<8} {minkowski(V[x], V[y]):.2f}")
    print()


if __name__ == "__main__":
    main()