"""Iteration: Short word policy — exact shape match for len<=3.

Both models agree: 1-3 char tokens are inherently ambiguous.
Require exact shape match or return 0.
Test impact on recall and false positive precision.
"""

import re
import unicodedata
from rapidfuzz import fuzz
from rapidfuzz.distance import JaroWinkler
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)

SHAPE = {}
for ch in "il1I|!tfj": SHAPE[ch] = "I"
for ch in "oO0cCeQa": SHAPE[ch] = "O"
for ch in "mwMW": SHAPE[ch] = "M"
for ch in "nuvNUV": SHAPE[ch] = "N"
for ch in "bhk6": SHAPE[ch] = "B"
for ch in "pqgy9": SHAPE[ch] = "P"
for ch in "dDPRB": SHAPE[ch] = "D"
for ch in "xXzZK": SHAPE[ch] = "X"
for ch in "sS5$": SHAPE[ch] = "S"
for ch in "-_~=": SHAPE[ch] = "H"
for ch in "234": SHAPE[ch] = "2"
for ch in "78": SHAPE[ch] = "7"
SHAPE["r"] = "R"

def shape_compress(t):
    t = unicodedata.normalize("NFKC", t).lower()
    return "".join(SHAPE.get(c, c) for c in t)

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)


# Using the best base from before: Max(JW, Shape_Lev) * len
MULTI_RULES = [("rn","m"),("nn","m"),("ri","n"),("cl","d"),("vv","w"),("ii","u")]
def shape_full(t):
    t = unicodedata.normalize("NFKC", t).lower()
    for p, r in MULTI_RULES:
        t = t.replace(p, r)
    return "".join(SHAPE.get(c, c) for c in t)


def score_short_guard_jw(query, candidate):
    """JW with short-word exact-shape policy."""
    q, c = norm(query), norm(candidate)
    if not q or not c:
        return 0.0
    if len(q) <= 3:
        return 0.95 if shape_compress(q) == shape_compress(c) else 0.0
    return JaroWinkler.similarity(q, c)


def score_short_guard_max(query, candidate):
    """Max(JW, Shape_Lev) * len with short-word policy."""
    q, c = norm(query), norm(candidate)
    if not q or not c:
        return 0.0
    if len(q) <= 3:
        return 0.95 if shape_compress(q) == shape_compress(c) else 0.0
    jw = JaroWinkler.similarity(q, c)
    shape_lev = fuzz.ratio(shape_full(q), shape_full(c)) / 100.0
    lr = length_ratio(q, c)
    return max(jw, shape_lev) * (0.8 + 0.2 * lr)


if __name__ == "__main__":
    evaluate(score_short_guard_jw, name="18a - JW + short guard (shape exact)")
    evaluate(score_short_guard_max, name="18b - Max(JW,Shape)*len + short guard")
