"""Iteration: Shape-compressed bigram Dice coefficient.

GPT's idea: generate bigrams from shape-compressed text, score with
Sorensen-Dice. Robust to severe garbling because even heavily corrupted
text preserves some local adjacency.
"""

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

MULTI_RULES = [("rn","m"),("nn","m"),("ri","n"),("cl","d"),("vv","w"),("ii","u")]
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()
    for p, r in MULTI_RULES:
        t = t.replace(p, r)
    return "".join(SHAPE.get(c, c) for c in t)


def bigrams(s):
    """Generate character bigrams from a string."""
    return set(s[i:i+2] for i in range(len(s) - 1))


def dice_coefficient(a, b):
    """Sorensen-Dice coefficient on two sets."""
    if not a and not b:
        return 1.0
    if not a or not b:
        return 0.0
    intersection = len(a & b)
    return 2.0 * intersection / (len(a) + len(b))


def trigrams(s):
    return set(s[i:i+3] for i in range(len(s) - 2))


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)


# --- Variants ---

def score_raw_bigram_dice(query, candidate):
    """Bigram Dice on raw (lowercased, punct-stripped) text."""
    q = strip_punct(query).lower()
    c = strip_punct(candidate).lower()
    if len(q) < 2 or len(c) < 2:
        return 1.0 if q == c else 0.0
    return dice_coefficient(bigrams(q), bigrams(c))


def score_shape_bigram_dice(query, candidate):
    """Bigram Dice on shape-compressed text."""
    q = shape_compress(strip_punct(query))
    c = shape_compress(strip_punct(candidate))
    if len(q) < 2 or len(c) < 2:
        return 1.0 if q == c else 0.0
    return dice_coefficient(bigrams(q), bigrams(c))


def score_shape_trigram_dice(query, candidate):
    """Trigram Dice on shape-compressed text."""
    q = shape_compress(strip_punct(query))
    c = shape_compress(strip_punct(candidate))
    if len(q) < 3 or len(c) < 3:
        return 1.0 if q == c else 0.0
    return dice_coefficient(trigrams(q), trigrams(c))


def score_shape_bigram_dice_len(query, candidate):
    """Shape bigram Dice * length penalty."""
    q = shape_compress(strip_punct(query))
    c = shape_compress(strip_punct(candidate))
    if len(q) < 2 or len(c) < 2:
        return 1.0 if q == c else 0.0
    d = dice_coefficient(bigrams(q), bigrams(c))
    lr = length_ratio(strip_punct(query), strip_punct(candidate))
    return d * (0.8 + 0.2 * lr)


if __name__ == "__main__":
    evaluate(score_raw_bigram_dice, name="16a - Raw Bigram Dice")
    evaluate(score_shape_bigram_dice, name="16b - Shape Bigram Dice")
    evaluate(score_shape_trigram_dice, name="16c - Shape Trigram Dice")
    evaluate(score_shape_bigram_dice_len, name="16d - Shape Bigram Dice * len")
