"""Two-stage matching: permissive recall stage → precision re-ranker.

Stage 1: max(JW, RevJW, Shape_Lev) — grab top-K candidates (permissive)
Stage 2: re-rank with blend of all features, including expensive ones

Tests several re-ranker designs against the false positive dataset.
"""

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)

def shape_simple(t):
    t = unicodedata.normalize("NFKC", t).lower()
    return "".join(SHAPE.get(c, c) for c in t)

# Weighted Levenshtein (expensive — only for re-ranker)
SHAPE_CLASS = {}
for cls, chars in [("V","il1I|!tfj"),("O","oO0cCeQa"),("M","mwMW"),("N","nuvNUV"),
                    ("B","bhk6"),("P","pqgy9"),("D","dDPRB"),("X","xXzZK"),("S","sS5$")]:
    for ch in chars: SHAPE_CLASS[ch.lower()] = cls

def weighted_lev_dist(s1, s2):
    n, m = len(s1), len(s2)
    dp = [[0.0]*(m+1) for _ in range(n+1)]
    for i in range(n+1): dp[i][0] = float(i)
    for j in range(m+1): dp[0][j] = float(j)
    for i in range(1, n+1):
        for j in range(1, m+1):
            if s1[i-1] == s2[j-1]:
                cost = 0.0
            elif SHAPE_CLASS.get(s1[i-1]) and SHAPE_CLASS.get(s1[i-1]) == SHAPE_CLASS.get(s2[j-1]):
                cost = 0.2
            else:
                cost = 1.0
            dp[i][j] = min(dp[i-1][j]+1, dp[i][j-1]+1, dp[i-1][j-1]+cost)
    return dp[n][m]

def weighted_lev_sim(a, b):
    if not a or not b: return 0.0
    if a == b: return 1.0
    d = weighted_lev_dist(a, b)
    return max(0.0, 1.0 - d / max(len(a), len(b)))


# ============================================================
# STAGE 1: Recall-optimized scorer (cheap, permissive)
# ============================================================

def stage1_score(query, candidate):
    """max(JW, RevJW, Shape_Lev) — no length penalty, no short guard.
    Purpose: never miss a true match."""
    q, c = norm(query), norm(candidate)
    if not q or not c: return 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
    return max(jw, rev, sl)


# ============================================================
# STAGE 2: Precision re-rankers (run on top-K from stage 1)
# ============================================================

def compute_features(query, candidate):
    """Compute all features for a (query, candidate) pair."""
    q, c = norm(query), norm(candidate)
    if not q or not c:
        return {"jw": 0, "rev_jw": 0, "shape_lev": 0, "weighted_lev": 0, "len_ratio": 0, "len_diff": 0}
    sq, sc = shape_compress(q), shape_compress(c)
    return {
        "jw": JaroWinkler.similarity(q, c),
        "rev_jw": JaroWinkler.similarity(q[::-1], c[::-1]),
        "shape_lev": fuzz.ratio(sq, sc) / 100.0,
        "weighted_lev": weighted_lev_sim(q, c),
        "len_ratio": length_ratio(q, c),
        "len_diff": abs(len(q) - len(c)),
    }


def rerank_A_blend(query, candidate):
    """Linear blend of all features."""
    f = compute_features(query, candidate)
    q = norm(query)
    if len(q) <= 3:
        return 0.95 if shape_simple(q) == shape_simple(norm(candidate)) else 0.0
    return (0.25 * f["jw"] + 0.10 * f["rev_jw"] +
            0.25 * f["shape_lev"] + 0.15 * f["weighted_lev"] +
            0.15 * f["len_ratio"] + 0.10 * max(0, 1.0 - f["len_diff"] / 5.0))


def rerank_B_gated_max(query, candidate):
    """Gated max: require corroboration when top score is borderline."""
    f = compute_features(query, candidate)
    q = norm(query)
    if len(q) <= 3:
        return 0.95 if shape_simple(q) == shape_simple(norm(candidate)) else 0.0
    scores = [f["jw"], f["rev_jw"], f["shape_lev"]]
    scores.sort(reverse=True)
    m = scores[0]      # best
    s2 = scores[1]     # second best
    lr = f["len_ratio"]

    # If top score is very high, trust it
    if m >= 0.92:
        return m * (0.85 + 0.15 * lr)
    # If second score corroborates, trust it
    if s2 >= 0.70:
        return m * (0.85 + 0.15 * lr)
    # Otherwise penalize — single weak signal is suspicious
    return m * (0.85 + 0.15 * lr) * 0.90


def rerank_C_power_mean(query, candidate):
    """Power mean (p=8): smooth max that slightly prefers corroboration."""
    f = compute_features(query, candidate)
    q = norm(query)
    if len(q) <= 3:
        return 0.95 if shape_simple(q) == shape_simple(norm(candidate)) else 0.0
    scores = [f["jw"], f["rev_jw"], f["shape_lev"]]
    p = 8
    pm = (sum(s**p for s in scores) / len(scores)) ** (1.0/p)
    lr = f["len_ratio"]
    return pm * (0.85 + 0.15 * lr)


def rerank_D_blend_with_weighted(query, candidate):
    """Blend that includes weighted Levenshtein (expensive, precision-focused)."""
    f = compute_features(query, candidate)
    q = norm(query)
    if len(q) <= 3:
        return 0.95 if shape_simple(q) == shape_simple(norm(candidate)) else 0.0
    return (0.20 * f["jw"] + 0.10 * f["rev_jw"] +
            0.25 * f["shape_lev"] + 0.20 * f["weighted_lev"] +
            0.15 * f["len_ratio"] + 0.10 * max(0, 1.0 - f["len_diff"] / 5.0))


def rerank_E_confidence_gap(query, candidate, _all_scores=None):
    """Blend + confidence gap bonus: if this candidate is far ahead of #2,
    boost its score. (Only usable in the two-stage context.)"""
    # This one is special — can't be used as a simple scorer.
    # We'll handle it differently in the evaluation loop.
    pass


RERANKERS = {
    "A_blend": rerank_A_blend,
    "B_gated_max": rerank_B_gated_max,
    "C_power_mean": rerank_C_power_mean,
    "D_blend_wtd_lev": rerank_D_blend_with_weighted,
}


# ============================================================
# Evaluation
# ============================================================

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 evaluate_two_stage(reranker_name, reranker_fn, top_k=10):
    """Two-stage evaluation:
    Stage 1: score all candidates with stage1_score, take top-K
    Stage 2: re-rank top-K with reranker_fn
    """
    pairs = load_pairs()
    rank_1 = rank_top3 = total = 0
    failures = []

    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

        # Stage 1: recall-optimized, get top-K
        s1_scored = [(stage1_score(correction, c), c) for c in candidates]
        s1_scored.sort(key=lambda x: x[0], reverse=True)
        top_candidates = [c for _, c in s1_scored[:top_k]]

        # Stage 2: re-rank top-K
        s2_scored = [(reranker_fn(correction, c), c) for c in top_candidates]
        s2_scored.sort(key=lambda x: x[0], reverse=True)

        # Find rank of correct answer
        rank = None
        for i, (s, cand) in enumerate(s2_scored):
            cc, ec = strip_punct(cand).lower(), strip_punct(ocr_error).lower()
            if ec in cc or cc in ec or cc == ec:
                rank = i + 1
                break

        if rank is None:
            # Check if it was filtered by stage 1
            s1_rank = None
            for i, (s, cand) in enumerate(s1_scored):
                cc, ec = strip_punct(cand).lower(), strip_punct(ocr_error).lower()
                if ec in cc or cc in ec or cc == ec:
                    s1_rank = i + 1
                    break
            if s1_rank and s1_rank > top_k:
                # Stage 1 filtered it out
                total += 1
                if len(failures) < 8:
                    failures.append((correction, ocr_error, f"s1>{top_k}", s1_scored[0][1][:25], s1_scored[0][0], 0))
                continue
            else:
                continue

        total += 1
        if rank == 1: rank_1 += 1
        if rank <= 3: rank_top3 += 1
        if rank != 1 and len(failures) < 8:
            failures.append((correction, ocr_error, rank, s2_scored[0][1][:25], s2_scored[0][0], s2_scored[rank-1][0] if rank <= len(s2_scored) else 0))

    print(f"\n  {reranker_name:<28} Rank1: {rank_1}/{total} ({rank_1/total:.1%})  Top3: {rank_top3}/{total} ({rank_top3/total:.1%})")
    for q, e, r, w, ws, cs in failures:
        print(f"    {q:<16} exp={e:<16} rank={str(r):<6} winner={w:<25} ({ws:.2f})")


def evaluate_recall_stage2(reranker_name, reranker_fn):
    """Test recall of the re-ranker on direct pairs."""
    from harness import evaluate
    evaluate(reranker_fn, name=f"Recall: {reranker_name}", show_failures=5)


if __name__ == "__main__":
    # First: recall on direct pairs (do re-rankers still score true pairs well?)
    print("=" * 70)
    print("  RECALL ON DIRECT PAIRS (re-ranker scores)")
    print("=" * 70)
    for name, fn in RERANKERS.items():
        evaluate_recall_stage2(name, fn)

    # Then: false positive precision with two-stage pipeline
    for k in [5, 10, 20]:
        print("\n" + "=" * 70)
        print(f"  TWO-STAGE FALSE POSITIVES (top-K={k})")
        print("=" * 70)

        # Baseline: stage 1 only (no re-rank)
        evaluate_two_stage(f"Stage1 only (K={k})", stage1_score, top_k=k)

        for name, fn in RERANKERS.items():
            evaluate_two_stage(f"{name} (K={k})", fn, top_k=k)
