"""Iteration: Weighted edit distance using shape classes as soft costs.

GPT's idea: instead of hard-collapsing to shape classes, use them as
cheaper substitution costs in Levenshtein.
- same char: 0
- same shape class: 0.2
- different: 1.0

Then normalize to [0,1] similarity.
"""

import re
import unicodedata
from harness import evaluate

PUNCT_RE = re.compile(r'^[^a-zA-Z0-9]+|[^a-zA-Z0-9]+$')

def strip_punct(t):
    return PUNCT_RE.sub('', t)

# Build shape class lookup
SHAPE_CLASS = {}
classes = [
    ("V", "il1I|!tfj"),
    ("O", "oO0cCeQa"),
    ("M", "mwMW"),
    ("N", "nuvNUV"),
    ("B", "bhk6"),
    ("P", "pqgy9"),
    ("D", "dDPRB"),
    ("X", "xXzZK"),
    ("S", "sS5$"),
]
for cls, chars in classes:
    for ch in chars:
        SHAPE_CLASS[ch.lower()] = cls


def sub_cost(a, b):
    if a == b:
        return 0.0
    ca = SHAPE_CLASS.get(a)
    cb = SHAPE_CLASS.get(b)
    if ca and cb and ca == cb:
        return 0.2
    return 1.0


def weighted_levenshtein(s1, s2):
    """Weighted Levenshtein with shape-class soft costs."""
    n, m = len(s1), len(s2)
    dp = [[0.0] * (m + 1) for _ in range(n + 1)]
    for i in range(n + 1):
        dp[i][0] = float(i)
    for j in range(m + 1):
        dp[0][j] = float(j)
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            cost = sub_cost(s1[i-1], s2[j-1])
            dp[i][j] = min(
                dp[i-1][j] + 1.0,      # deletion
                dp[i][j-1] + 1.0,      # insertion
                dp[i-1][j-1] + cost,    # substitution
            )
    return dp[n][m]


def norm(t):
    return strip_punct(unicodedata.normalize("NFKC", t)).lower()

def length_ratio(a, b):
    la, lb = len(a), len(b)
    if la == 0 or lb == 0:
        return 0.0
    return min(la, lb) / max(la, lb)


def score_weighted_lev(query, candidate):
    """Weighted Levenshtein with shape soft costs, normalized to [0,1]."""
    q, c = norm(query), norm(candidate)
    if not q or not c:
        return 0.0
    if q == c:
        return 1.0
    dist = weighted_levenshtein(q, c)
    max_len = max(len(q), len(c))
    return max(0.0, 1.0 - dist / max_len)


def score_weighted_lev_len(query, candidate):
    """Weighted Lev + length penalty."""
    q, c = norm(query), norm(candidate)
    if not q or not c:
        return 0.0
    if q == c:
        return 1.0
    dist = weighted_levenshtein(q, c)
    max_len = max(len(q), len(c))
    sim = max(0.0, 1.0 - dist / max_len)
    lr = length_ratio(q, c)
    return sim * (0.8 + 0.2 * lr)


if __name__ == "__main__":
    evaluate(score_weighted_lev, name="19a - Weighted Lev (shape soft costs)")
    evaluate(score_weighted_lev_len, name="19b - Weighted Lev * len penalty")
