"""Side-by-side failure analysis: what does each algorithm get WRONG?

For each algorithm, show cases where it fails (rank != 1) and what it
picked instead. Focus on cases where the algorithms DISAGREE — one gets
it right and the other doesn't.
"""

import csv
import re
import unicodedata
from pathlib import Path

from rapidfuzz import fuzz
from rapidfuzz.distance import JaroWinkler

DATA_DIR = Path(__file__).parent
CSV_PATH = DATA_DIR / "ocr_errors.csv"
OUTPUT_DIR = DATA_DIR / "output"
PUNCT_RE = re.compile(r'^[^a-zA-Z0-9]+|[^a-zA-Z0-9]+$')

def strip_punct(t): return PUNCT_RE.sub('', t)
def norm(t): return strip_punct(unicodedata.normalize("NFKC", t)).lower()
def length_ratio(a, b):
    la, lb = len(a), len(b)
    return min(la, lb) / max(la, lb) if la and lb else 0.0

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)


# The two contenders
def score_blend(query, candidate):
    """Best precision: .4*JW + .4*Shape + .2*LenRatio (from script 14)."""
    q, c = norm(query), norm(candidate)
    if not q or not c: return 0.0
    jw = JaroWinkler.similarity(q, c)
    sq, sc = shape_compress(q), shape_compress(c)
    sl = fuzz.ratio(sq, sc) / 100.0
    lr = length_ratio(q, c)
    return 0.4 * jw + 0.4 * sl + 0.2 * lr


def score_max_recall(query, candidate):
    """Best recall: Max(JW, RevJW, Shape)*len + short guard (20A)."""
    q, c = norm(query), norm(candidate)
    if not q or not c: return 0.0
    if len(q) <= 3:
        sq = unicodedata.normalize("NFKC", q).lower()
        sc = unicodedata.normalize("NFKC", c).lower()
        sq = "".join(SHAPE.get(ch, ch) for ch in sq)
        sc = "".join(SHAPE.get(ch, ch) for ch in sc)
        return 0.95 if sq == sc else 0.0
    jw = JaroWinkler.similarity(q, c)
    rev = JaroWinkler.similarity(q[::-1], c[::-1])
    sq, sc = shape_compress(q), shape_compress(c)
    sl = fuzz.ratio(sq, sc) / 100.0
    lr = length_ratio(q, c)
    return max(jw, rev, sl) * (0.8 + 0.2 * lr)


def load_pairs():
    pairs = []
    with open(CSV_PATH) as f:
        for row in csv.DictReader(f): pairs.append((row["ocr_error"], row["correction"]))
    return pairs

def find_source_file(ocr_error):
    for p in OUTPUT_DIR.glob("*.txt"):
        if ocr_error in p.read_text(errors="replace"): return p
    return None

def get_candidates(text, n_words):
    tokens = text.split()
    if n_words == 1: return tokens
    return [" ".join(tokens[i:i+n_words]) for i in range(len(tokens) - n_words + 1)]


def get_rank(scored_list, ocr_error):
    for i, (s, cand) in enumerate(scored_list):
        cc, ec = strip_punct(cand).lower(), strip_punct(ocr_error).lower()
        if ec in cc or cc in ec or cc == ec:
            return i + 1, s
    return None, None


if __name__ == "__main__":
    pairs = load_pairs()

    blend_only_wrong = []   # blend wrong, max right
    max_only_wrong = []     # max wrong, blend right
    both_wrong = []         # both wrong
    both_right = 0
    total = 0

    for ocr_error, correction in pairs:
        src = find_source_file(ocr_error)
        if not src: continue
        text = src.read_text(errors="replace")
        n_words = len(correction.split())
        candidates = [c for c in get_candidates(text, n_words)
                      if strip_punct(c).lower() != strip_punct(correction).lower()]
        if not candidates: continue

        # Score with both
        blend_scored = sorted([(score_blend(correction, c), c) for c in candidates],
                              key=lambda x: x[0], reverse=True)
        max_scored = sorted([(score_max_recall(correction, c), c) for c in candidates],
                            key=lambda x: x[0], reverse=True)

        blend_rank, blend_corr_score = get_rank(blend_scored, ocr_error)
        max_rank, max_corr_score = get_rank(max_scored, ocr_error)

        if blend_rank is None or max_rank is None:
            continue

        total += 1
        blend_ok = blend_rank == 1
        max_ok = max_rank == 1

        entry = {
            "query": correction,
            "expected": ocr_error,
            "blend_rank": blend_rank,
            "blend_winner": blend_scored[0][1][:25],
            "blend_winner_score": blend_scored[0][0],
            "blend_corr_score": blend_corr_score,
            "max_rank": max_rank,
            "max_winner": max_scored[0][1][:25],
            "max_winner_score": max_scored[0][0],
            "max_corr_score": max_corr_score,
        }

        if blend_ok and max_ok:
            both_right += 1
        elif blend_ok and not max_ok:
            max_only_wrong.append(entry)
        elif not blend_ok and max_ok:
            blend_only_wrong.append(entry)
        else:
            both_wrong.append(entry)

    print(f"Total tested: {total}")
    print(f"Both right:   {both_right} ({both_right/total:.1%})")
    print(f"Blend right, Max wrong: {len(max_only_wrong)} ({len(max_only_wrong)/total:.1%})")
    print(f"Max right, Blend wrong: {len(blend_only_wrong)} ({len(blend_only_wrong)/total:.1%})")
    print(f"Both wrong:   {len(both_wrong)} ({len(both_wrong)/total:.1%})")

    def show(title, entries, limit=25):
        print(f"\n{'='*100}")
        print(f"  {title} ({len(entries)} cases)")
        print(f"{'='*100}")
        print(f"  {'Query':<15} {'Expected':<15} {'B_rk':>4} {'B_winner':<22} {'B_ws':>5} {'B_cs':>5}  {'M_rk':>4} {'M_winner':<22} {'M_ws':>5} {'M_cs':>5}")
        print(f"  {'-'*100}")
        for e in entries[:limit]:
            print(f"  {e['query']:<15} {e['expected']:<15} {e['blend_rank']:>4} {e['blend_winner']:<22} {e['blend_winner_score']:>5.2f} {e['blend_corr_score']:>5.02f}  {e['max_rank']:>4} {e['max_winner']:<22} {e['max_winner_score']:>5.02f} {e['max_corr_score']:>5.02f}")

    show("BLEND RIGHT, MAX WRONG (blend's unique wins)", max_only_wrong)
    show("MAX RIGHT, BLEND WRONG (max's unique wins)", blend_only_wrong)
    show("BOTH WRONG (irreducible failures)", both_wrong)
