import random
import unittest

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]


class TestLevenshtein(unittest.TestCase):
    def test_known(self):
        cases = [("kitten", "sitting", 3), ("", "abc", 3), ("abc", "", 3),
                 ("flaw", "lawn", 2), ("same", "same", 0), ("a", "b", 1),
                 ("", "", 0)]
        for a, b, e in cases:
            with self.subTest(a=a, b=b):
                self.assertEqual(edit_distance(a, b), e)

    def test_random_vs_brute(self):
        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))
            with self.subTest(a=a, b=b):
                self.assertEqual(edit_distance(a, b), brute(a, b))

    def test_ops_roundtrip(self):
        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)
            out = []
            i = 0
            j = 0
            for tag in o:
                if tag == "=":
                    self.assertEqual(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":
                    out.append(b[j])
                    j += 1
            self.assertEqual("".join(out), b, (a, b, o, "".join(out)))
            edits = sum(1 for t in o if t in ("S", "D", "I"))
            self.assertEqual(edits, edit_distance(a, b), (a, b, edits, edit_distance(a, b)))


if __name__ == "__main__":
    unittest.main()