"""Iteration B: Length-aware scoring — penalize candidates whose length
differs significantly from the query.

Hypothesis: many false positives come from short tokens matching queries
(e.g., "be" matching "began", "W" matching "was"). Adding a length ratio
penalty should suppress these without hurting true matches (which are
typically similar length to the query).
"""

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(token):
    return PUNCT_RE.sub('', token)


def length_ratio(a, b):
    """Returns ratio of shorter to longer length. 1.0 = same length."""
    la, lb = len(a), len(b)
    if la == 0 or lb == 0:
        return 0.0
    return min(la, lb) / max(la, lb)


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)]


# JW + punct strip + length penalty
def score_jw_len(query, candidate):
    q = strip_punct(query).lower()
    c = strip_punct(candidate).lower()
    if not q or not c:
        return 0.0
    sim = JaroWinkler.similarity(q, c)
    lr = length_ratio(q, c)
    # Blend: mostly similarity, but penalize length mismatch
    return sim * (0.7 + 0.3 * lr)

# Lev + punct strip + length penalty
def score_lev_len(query, candidate):
    q = strip_punct(query).lower()
    c = strip_punct(candidate).lower()
    if not q or not c:
        return 0.0
    sim = fuzz.ratio(q, c) / 100.0
    lr = length_ratio(q, c)
    return sim * (0.7 + 0.3 * lr)

# Shape + punct strip + length penalty
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_len(query, candidate):
    q = strip_punct(query)
    c = strip_punct(candidate)
    if not q or not c:
        return 0.0
    sim = fuzz.ratio(_compress(q), _compress(c)) / 100.0
    lr = length_ratio(q, c)
    return sim * (0.7 + 0.3 * lr)


ALGORITHMS = {
    "JW+len": score_jw_len,
    "Lev+len": score_lev_len,
    "Shape+len": score_shape_len,
}


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)
        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
        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[:28], "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__":
    # Also run recall test on the pairs
    from harness import evaluate as eval_recall

    print("=== RECALL (pair matching) ===")
    for name, fn in ALGORITHMS.items():
        eval_recall(fn, name=name, show_failures=5)

    print("\n=== FALSE POSITIVE PRECISION ===")
    for name, fn in ALGORITHMS.items():
        evaluate_fp(name, fn)
