"""Iteration C: Two-stage hybrid — shape compression as fast filter,
then precise JW scoring on survivors.

Hypothesis: shape compression has high recall (catches the right answer)
but creates collisions. JW has better discrimination. Use shape as a
gate to get candidates, then re-rank with JW.

Also tests: combining multiple signals (shape score + JW score + length ratio)
as a weighted blend.
"""

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):
    la, lb = len(a), len(b)
    if la == 0 or lb == 0:
        return 0.0
    return min(la, lb) / max(la, lb)

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)


# --- Hybrid scoring functions ---

def score_shape_then_jw(query, candidate):
    """Two-stage: shape filter (must be >= 0.5) then JW re-rank."""
    q = strip_punct(query)
    c = strip_punct(candidate)
    if not q or not c:
        return 0.0
    shape_sim = fuzz.ratio(_compress(q), _compress(c)) / 100.0
    if shape_sim < 0.5:
        return 0.0  # filtered out
    jw = JaroWinkler.similarity(q.lower(), c.lower())
    return jw


def score_blend(query, candidate):
    """Weighted blend of JW + shape + length ratio."""
    q = strip_punct(query)
    c = strip_punct(candidate)
    if not q or not c:
        return 0.0
    jw = JaroWinkler.similarity(q.lower(), c.lower())
    shape_sim = fuzz.ratio(_compress(q), _compress(c)) / 100.0
    lr = length_ratio(q, c)
    # Blend weights: JW contributes precision, shape contributes recall, length prevents short-match junk
    return 0.4 * jw + 0.4 * shape_sim + 0.2 * lr


def score_max_jw_shape(query, candidate):
    """Take the max of JW and shape — whichever is more confident."""
    q = strip_punct(query)
    c = strip_punct(candidate)
    if not q or not c:
        return 0.0
    jw = JaroWinkler.similarity(q.lower(), c.lower())
    shape_sim = fuzz.ratio(_compress(q), _compress(c)) / 100.0
    lr = length_ratio(q, c)
    return max(jw, shape_sim) * (0.8 + 0.2 * lr)


def score_blend_v2(query, candidate):
    """JW * shape * length — multiplicative blend.
    All three must be decent for a high score."""
    q = strip_punct(query)
    c = strip_punct(candidate)
    if not q or not c:
        return 0.0
    jw = JaroWinkler.similarity(q.lower(), c.lower())
    shape_sim = fuzz.ratio(_compress(q), _compress(c)) / 100.0
    lr = length_ratio(q, c)
    # Geometric-ish mean: all three must be high
    return (jw * shape_sim) ** 0.5 * (0.8 + 0.2 * lr)


ALGORITHMS = {
    "ShapeThenJW": score_shape_then_jw,
    "Blend(.4/.4/.2)": score_blend,
    "Max(JW,Shape)+len": score_max_jw_shape,
    "Blend_v2(mult)": score_blend_v2,
}


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


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) < 12:
            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.02f} {f['correct_score']:>5.02f}")
    print()


if __name__ == "__main__":
    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)
