"""Minimum edit distance, worked out the way Chapter 2 works it out.

Run this and you get the same three tables and the same trace the book prints
for `cat` to `cart`. Change the two words at the bottom and it works them out
for yours.

    python3 edit_distance.py                  # cat -> cart, as in the book
    python3 edit_distance.py sunday saturday  # the chapter's exercise
    python3 edit_distance.py cat cut --sub 2  # see the penalties bite

The code is written to be read, not to be fast. Every step of the recurrence
is spelled out.

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

import argparse


def costs(sub_cost=1, ins_cost=1, del_cost=1):
    """The penalties. Levenshtein uses 1, 1, and 0-or-1 for a substitution.

    Jurafsky and Martin often use sub_cost=2, reasoning that a substitution is
    really a deletion plus an insertion. Try both and watch the answer move.
    """
    def sub(a, b):
        return 0 if a == b else sub_cost      # a match is free
    return sub, (lambda b: ins_cost), (lambda a: del_cost)


def build_table(s, t, sub_cost=1, ins_cost=1, del_cost=1):
    """Return the full (len(s)+1) x (len(t)+1) table of distances."""
    sub, ins, dele = costs(sub_cost, ins_cost, del_cost)
    D = [[0] * (len(t) + 1) for _ in range(len(s) + 1)]

    # Base cases: turning a prefix into the empty string costs one edit per
    # character, and the other way round costs one insertion per character.
    for i in range(len(s) + 1):
        D[i][0] = i * del_cost
    for j in range(len(t) + 1):
        D[0][j] = j * ins_cost

    for i in range(1, len(s) + 1):
        for j in range(1, len(t) + 1):
            delete = D[i - 1][j] + dele(s[i - 1])
            insert = D[i][j - 1] + ins(t[j - 1])
            match = D[i - 1][j - 1] + sub(s[i - 1], t[j - 1])
            D[i][j] = min(delete, insert, match)
    return D


def trace_back(D, s, t, sub_cost=1, ins_cost=1, del_cost=1):
    """Walk from the answer cell to the origin, recovering the edits.

    At each cell we ask which of the three neighbours actually produced this
    value. Diagonal is preferred, so a match is reported as a match rather
    than as a delete plus an insert that happens to cost the same.
    """
    sub, ins, dele = costs(sub_cost, ins_cost, del_cost)
    i, j = len(s), len(t)
    path, steps = [(i, j)], []
    while i > 0 or j > 0:
        if i > 0 and j > 0 and D[i][j] == D[i-1][j-1] + sub(s[i-1], t[j-1]):
            op = "match" if s[i-1] == t[j-1] else f"substitute {s[i-1]}->{t[j-1]}"
            frm, i, j = (i-1, j-1), i-1, j-1
        elif j > 0 and D[i][j] == D[i][j-1] + ins(t[j-1]):
            op, frm, j = f"insert {t[j-1]}", (i, j-1), j - 1
        else:
            op, frm, i = f"delete {s[i-1]}", (i-1, j), i - 1
        steps.append((path[-1], frm, op))
        path.append(frm)
    return list(reversed(path)), list(reversed(steps))


def show(D, s, t, rows=None, cols=None, path=None, box=False):
    """Print the table, optionally only the top-left corner, path cells starred."""
    rows = len(s) + 1 if rows is None else rows
    cols = len(t) + 1 if cols is None else cols
    path = set(path or [])
    head = "    " + "".join(f"{c:>5}" for c in ["#"] + list(t)[:cols - 1])
    print(head)
    print("    " + "-" * (5 * cols))
    for i in range(rows):
        label = "#" if i == 0 else s[i - 1]
        line = f"{label:>3} |"
        for j in range(cols):
            cell = str(D[i][j])
            if box and (i, j) == (len(s), len(t)):
                cell = f"[{cell}]"                 # the answer
            elif (i, j) in path:
                cell = f"*{cell}"                  # on the cheapest path
            line += f"{cell:>5}"
        print(line)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("source", nargs="?", default="cat")
    ap.add_argument("target", nargs="?", default="cart")
    ap.add_argument("--sub", type=int, default=1, help="substitution penalty")
    ap.add_argument("--ins", type=int, default=1, help="insertion penalty")
    ap.add_argument("--del", dest="dele", type=int, default=1,
                    help="deletion penalty")
    a = ap.parse_args()
    s, t = a.source, a.target

    print(f"\n{s!r} -> {t!r}   penalties: sub={a.sub} ins={a.ins} del={a.dele}")
    D = build_table(s, t, a.sub, a.ins, a.dele)

    print("\nSTEP 1: the first four cells")
    print("  base cases fill the top row and the left column;")
    print(f"  D[1,1] compares {s[0]!r} with {t[0]!r}:")
    print(f"    delete -> D[0,1] + {a.dele} = {D[0][1] + a.dele}")
    print(f"    insert -> D[1,0] + {a.ins} = {D[1][0] + a.ins}")
    same = s[0] == t[0]
    print(f"    {'match' if same else 'substitute'}  -> D[0,0] + "
          f"{0 if same else a.sub} = {0 if same else a.sub}")
    print()
    show(D, s, t, rows=2, cols=2)

    print("\nSTEP 2: the first pass")
    show(D, s, t, rows=2)

    print("\nSTEP 3: the finished table (* = on the cheapest path, [] = answer)")
    path, steps = trace_back(D, s, t, a.sub, a.ins, a.dele)
    show(D, s, t, path=path, box=True)

    print(f"\nEDIT DISTANCE = {D[len(s)][len(t)]}\n")
    print("THE TRACE (read downwards to replay the edits):")
    print(f"  {'cell':<10}{'from':<10}{'operation'}")
    for (ci, cj), (fi, fj), op in steps:
        print(f"  D[{ci},{cj}]{'':<4}D[{fi},{fj}]{'':<4}{op}")
    kept = [op for _, _, op in steps if op != "match"]
    print(f"\n  {len(kept)} non-free edit(s): {', '.join(kept) or 'none'}")


if __name__ == "__main__":
    main()