"""Iteration A: Strip punctuation from tokens before scoring.

Hypothesis: most false positive "winners" are the correct word with attached
punctuation (e.g., "residence." beating "reaidence"). Stripping punctuation
should dramatically improve rank-1 precision across all algorithms.
"""

import csv
import re
import sys
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"

# Punctuation to strip from token edges
PUNCT_RE = re.compile(r'^[^a-zA-Z0-9]+|[^a-zA-Z0-9]+$')


def strip_punct(token):
    """Strip leading/trailing punctuation from a token."""
    return PUNCT_RE.sub('', token)


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 txt_path in OUTPUT_DIR.glob("*.txt"):
        content = txt_path.read_text(errors="replace")
        if ocr_error in content:
            return txt_path
    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)]


# --- Algorithms with punct stripping ---

def score_jw_clean(query, candidate):
    q = strip_punct(query).lower()
    c = strip_punct(candidate).lower()
    if not q or not c:
        return 0.0
    return JaroWinkler.similarity(q, c)

def score_lev_clean(query, candidate):
    q = strip_punct(query).lower()
    c = strip_punct(candidate).lower()
    if not q or not c:
        return 0.0
    return fuzz.ratio(q, c) / 100.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 _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 score_shape_clean(query, candidate):
    q = strip_punct(query)
    c = strip_punct(candidate)
    if not q or not c:
        return 0.0
    return fuzz.ratio(_compress(q), _compress(c)) / 100.0

def score_word_jw_clean(query, candidate):
    q_words = [strip_punct(w) for w in query.lower().split() if strip_punct(w)]
    c_words = [strip_punct(w) for w in candidate.lower().split() if strip_punct(w)]
    if not q_words or not c_words:
        return 0.0
    total = 0.0
    for qw in q_words:
        best = max(JaroWinkler.similarity(qw, cw) for cw in c_words)
        total += best
    return total / len(q_words)


ALGORITHMS = {
    "JW+strip": score_jw_clean,
    "Lev+strip": score_lev_clean,
    "Shape+strip": score_shape_clean,
    "WordJW+strip": score_word_jw_clean,
}


def evaluate_fp(algo_name, score_fn):
    pairs = load_pairs()
    rank_1 = 0
    rank_top3 = 0
    total_tested = 0
    false_pos = []

    for ocr_error, correction in pairs:
        src = find_source_file(ocr_error)
        if src is None:
            continue

        text = src.read_text(errors="replace")
        n_words = len(correction.split())
        all_candidates = get_candidates(text, n_words)

        # Remove exact matches (case-insensitive, after punct strip)
        candidates = [c for c in all_candidates
                      if strip_punct(c).lower() != strip_punct(correction).lower()]

        if not candidates:
            continue

        scored = [(score_fn(correction, c), c) for c in candidates]
        scored.sort(key=lambda x: x[0], reverse=True)

        rank = None
        correct_score = None
        for i, (s, cand) in enumerate(scored):
            cand_clean = strip_punct(cand).lower()
            err_clean = strip_punct(ocr_error).lower()
            if err_clean in cand_clean or cand_clean in err_clean or cand_clean == err_clean:
                rank = i + 1
                correct_score = s
                break

        if rank is None:
            continue

        total_tested += 1
        if rank == 1:
            rank_1 += 1
        if rank <= 3:
            rank_top3 += 1

        if rank != 1 and len(false_pos) < 15:
            top_s, top_c = scored[0]
            false_pos.append({
                "query": correction,
                "expected": ocr_error,
                "rank": rank,
                "top_match": top_c[:30],
                "top_score": top_s,
                "correct_score": correct_score,
            })

    print(f"\n{'=' * 70}")
    print(f"  {algo_name}")
    print(f"{'=' * 70}")
    print(f"  Tested: {total_tested}  |  Rank 1: {rank_1} ({rank_1/total_tested:.1%})  |  Top 3: {rank_top3} ({rank_top3/total_tested:.1%})")

    if false_pos:
        print(f"\n  {'Query':<18} {'Expected':<18} {'Rank':>4} {'Winner':<28} {'Win':>5} {'Corr':>5}")
        print(f"  {'-'*95}")
        for f in false_pos:
            print(f"  {f['query']:<18} {f['expected']:<18} {f['rank']:>4} {f['top_match']:<28} {f['top_score']:>5.2f} {f['correct_score']:>5.2f}")
    print()


if __name__ == "__main__":
    for name, fn in ALGORITHMS.items():
        evaluate_fp(name, fn)
