import random, string
from levenshtein import edit_distance, ops


def brute(a, b):
    n, m = len(a), len(b)
    dp = [[0]*(m+1) for _ in range(n+1)]
    for i in range(n+1): dp[i][0]=i
    for j in range(m+1): dp[0][j]=j
    for i in range(1,n+1):
        for j in range(1,m+1):
            c = 0 if a[i-1]==b[j-1] else 1
            dp[i][j]=min(dp[i-1][j]+1, dp[i][j-1]+1, dp[i-1][j-1]+c)
    return dp[n][m]


def test_known():
    cases = [("kitten","sitting",3),("","abc",3),("abc","",3),
             ("flaw","lawn",2),("same","same",0),("a","b",1),
             ("","",-0 if False else 0)]
    for a,b,e in cases:
        assert edit_distance(a,b)==e, (a,b,edit_distance(a,b),e)


def test_random_vs_brute():
    rng = random.Random(42)
    alpha = "ab"
    for _ in range(2000):
        la = rng.randint(0,8); lb = rng.randint(0,8)
        a = "".join(rng.choice(alpha) for _ in range(la))
        b = "".join(rng.choice(alpha) for _ in range(lb))
        assert edit_distance(a,b)==brute(a,b), (a,b)


def test_ops_roundtrip():
    rng = random.Random(7)
    for _ in range(500):
        a = "".join(rng.choice("ab") for _ in range(rng.randint(0,6)))
        b = "".join(rng.choice("ab") for _ in range(rng.randint(0,6)))
        o = ops(a,b)
        # apply ops to a -> should yield b
        out = []
        i = 0  # index into a
        j = 0  # index into b
        for tag in o:
            if tag == "=":
                assert a[i] == b[j]
                out.append(a[i]); i += 1; j += 1
            elif tag == "S":
                out.append(b[j]); i += 1; j += 1
            elif tag == "D":
                i += 1
            elif tag == "I":
                # need current b char; find next unconsumed b
                out.append(b[j]); j += 1
        assert "".join(out) == b, (a,b,o,"".join(out))
        edits = sum(1 for t in o if t in ("S","D","I"))
        assert edits == edit_distance(a,b), (a,b,edits,edit_distance(a,b))


if __name__ == "__main__":
    test_known(); test_random_vs_brute(); test_ops_roundtrip()
    print("ALL OK")