"""False positive evaluation: search for correction in source document text,
check that the OCR error ranks #1 among all candidates.

For each (ocr_error, correction) pair:
1. Find the .txt file containing ocr_error
2. Split text on whitespace into tokens
3. For multi-word corrections, create n-grams of matching word count
4. Score correction against every token/n-gram
5. Report: rank of the correct match, and what beat it
"""

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


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):
    """Find the .txt file containing the 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):
    """Split text into whitespace tokens. If n_words > 1, create n-grams."""
    tokens = text.split()
    if n_words == 1:
        return tokens
    # Create n-grams
    candidates = []
    for i in range(len(tokens) - n_words + 1):
        candidates.append(" ".join(tokens[i : i + n_words]))
    return candidates


# --- Scoring functions to test ---

def score_jw(query, candidate):
    return JaroWinkler.similarity(query.lower(), candidate.lower())

def score_lev(query, candidate):
    return fuzz.ratio(query.lower(), candidate.lower()) / 100.0

def score_shape(query, candidate):
    """Shape compression + Levenshtein."""
    import unicodedata
    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)

    return fuzz.ratio(compress(query), compress(candidate)) / 100.0

def score_word_jw(query, candidate):
    """Word-level JW."""
    q_words = query.lower().split()
    c_words = candidate.lower().split()
    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": score_jw,
    "Lev": score_lev,
    "Shape": score_shape,
    "WordJW": score_word_jw,
}


def evaluate_false_positives(algo_name, score_fn, max_pairs=None):
    """Run false positive evaluation for a scoring function."""
    pairs = load_pairs()
    if max_pairs:
        pairs = pairs[:max_pairs]

    rank_1 = 0
    rank_top3 = 0
    total_tested = 0
    failures = []

    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())
        candidates = get_candidates(text, n_words)

        if not candidates:
            continue

        # Score all candidates
        scored = []
        for cand in candidates:
            s = score_fn(correction, cand)
            scored.append((s, cand))

        scored.sort(key=lambda x: x[0], reverse=True)

        # Find rank of the correct answer
        # The OCR error might appear as part of a larger token for multi-word,
        # so check if the candidate contains the OCR error
        rank = None
        for i, (s, cand) in enumerate(scored):
            if ocr_error in cand or cand in ocr_error:
                rank = i + 1
                break

        total_tested += 1
        if rank == 1:
            rank_1 += 1
        if rank is not None and rank <= 3:
            rank_top3 += 1

        if rank != 1 and len(failures) < 15:
            top_match = scored[0] if scored else (0, "")
            failures.append({
                "correction": correction,
                "ocr_error": ocr_error,
                "rank": rank if rank else ">50",
                "top_score": top_match[0],
                "top_match": top_match[1][:30],
                "correct_score": next((s for s, c in scored if ocr_error in c or c in ocr_error), 0),
            })

    # Report
    print(f"\n{'=' * 70}")
    print(f"  FALSE POSITIVE TEST: {algo_name}")
    print(f"{'=' * 70}")
    print(f"  Pairs tested: {total_tested}")
    print(f"  Rank 1 (correct answer is top match): {rank_1}/{total_tested} = {rank_1/total_tested:.1%}")
    print(f"  Top 3 (correct in top 3):             {rank_top3}/{total_tested} = {rank_top3/total_tested:.1%}")

    if failures:
        print(f"\n  Failures (correct answer not rank 1):")
        print(f"  {'Query':<20} {'Expected':<20} {'Rank':>5} {'Top Match':<30} {'Top':>5} {'Corr':>5}")
        print(f"  {'-'*20} {'-'*20} {'-'*5} {'-'*30} {'-'*5} {'-'*5}")
        for f in failures[:15]:
            print(f"  {f['correction']:<20} {f['ocr_error']:<20} {str(f['rank']):>5} {f['top_match']:<30} {f['top_score']:>5.2f} {f['correct_score']:>5.2f}")
    print()


if __name__ == "__main__":
    algo = sys.argv[1] if len(sys.argv) > 1 else "all"

    if algo == "all":
        for name, fn in ALGORITHMS.items():
            evaluate_false_positives(name, fn)
    elif algo in ALGORITHMS:
        evaluate_false_positives(algo, ALGORITHMS[algo])
    else:
        print(f"Unknown algorithm: {algo}. Choose from: {list(ALGORITHMS.keys())}")
