"""Find the elbow: at what threshold do we get the best recall/precision tradeoff?

Look at the score distributions for:
1. True pairs (OCR error vs correction) — should score HIGH
2. Random non-matching pairs from the same documents — should score LOW

The elbow is where the two distributions separate cleanly.
"""

import csv
import random
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 score_blend(query, candidate):
    q, c = norm(query), norm(candidate)
    if not q or not c: return 0.0
    jw = JaroWinkler.similarity(q, c)
    sq, sc = shape_compress(q), shape_compress(c)
    sl = fuzz.ratio(sq, sc) / 100.0
    lr = length_ratio(q, c)
    return 0.4 * jw + 0.4 * sl + 0.2 * lr


def score_max(query, candidate):
    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
    lr = length_ratio(q, c)
    return max(jw, rev, sl) * (0.8 + 0.2 * lr)


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


if __name__ == "__main__":
    pairs = load_pairs()
    random.seed(42)

    # Collect scores for true pairs and random negative pairs
    true_scores_blend = []
    true_scores_max = []
    neg_scores_blend = []
    neg_scores_max = []

    # Build a pool of all tokens from all source files
    all_tokens = set()
    for p in OUTPUT_DIR.glob("*.txt"):
        text = p.read_text(errors="replace")
        for tok in text.split():
            cleaned = strip_punct(tok)
            if len(cleaned) >= 3:
                all_tokens.add(cleaned)
    all_tokens = list(all_tokens)

    # True pair scores
    for ocr_error, correction in pairs:
        true_scores_blend.append(score_blend(correction, ocr_error))
        true_scores_max.append(score_max(correction, ocr_error))

    # Negative pair scores: for each correction, score against random tokens
    for _, correction in pairs:
        # Pick 5 random tokens that are NOT the correction or its OCR variant
        negs = random.sample(all_tokens, min(5, len(all_tokens)))
        for neg in negs:
            if norm(neg) == norm(correction):
                continue
            neg_scores_blend.append(score_blend(correction, neg))
            neg_scores_max.append(score_max(correction, neg))

    # Distribution analysis
    thresholds = [0.50, 0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85, 0.90, 0.95]

    print("=" * 80)
    print("  ELBOW ANALYSIS: True pairs vs Random negatives")
    print("=" * 80)

    for name, true_s, neg_s in [("Blend", true_scores_blend, neg_scores_blend),
                                 ("Max", true_scores_max, neg_scores_max)]:
        print(f"\n  --- {name} ---")
        print(f"  True pairs:  n={len(true_s)}, mean={sum(true_s)/len(true_s):.3f}, "
              f"median={sorted(true_s)[len(true_s)//2]:.3f}")
        print(f"  Neg pairs:   n={len(neg_s)}, mean={sum(neg_s)/len(neg_s):.3f}, "
              f"median={sorted(neg_s)[len(neg_s)//2]:.3f}")

        print(f"\n  {'Thresh':>7} {'True>=':>8} {'Neg>=':>8} {'TPR':>8} {'FPR':>8} {'F1':>8} {'TPR-FPR':>8}")
        print(f"  {'-'*7} {'-'*8} {'-'*8} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")
        best_f1 = 0
        best_thresh = 0
        best_gap = 0
        best_gap_thresh = 0
        for t in thresholds:
            tp = sum(1 for s in true_s if s >= t)
            fp = sum(1 for s in neg_s if s >= t)
            tpr = tp / len(true_s)
            fpr = fp / len(neg_s)
            precision = tp / (tp + fp) if (tp + fp) > 0 else 0
            f1 = 2 * precision * tpr / (precision + tpr) if (precision + tpr) > 0 else 0
            gap = tpr - fpr
            print(f"  {t:>7.2f} {tp:>8} {fp:>8} {tpr:>8.1%} {fpr:>8.1%} {f1:>8.3f} {gap:>+8.3f}")
            if f1 > best_f1:
                best_f1 = f1
                best_thresh = t
            if gap > best_gap:
                best_gap = gap
                best_gap_thresh = t
        print(f"\n  Best F1: {best_f1:.3f} at threshold {best_thresh}")
        print(f"  Best TPR-FPR gap: {best_gap:+.3f} at threshold {best_gap_thresh}")

    # Histogram view
    print("\n" + "=" * 80)
    print("  SCORE DISTRIBUTION (Blend)")
    print("=" * 80)
    buckets = [(i/20, (i+1)/20) for i in range(20)]
    print(f"  {'Range':>12} {'True':>8} {'Neg':>8} {'Ratio':>8}")
    for lo, hi in buckets:
        tc = sum(1 for s in true_scores_blend if lo <= s < hi)
        nc = sum(1 for s in neg_scores_blend if lo <= s < hi)
        ratio = tc / nc if nc > 0 else float('inf')
        bar_t = '#' * min(tc // 2, 40)
        bar_n = '.' * min(nc // 50, 40)
        print(f"  [{lo:.2f},{hi:.2f}) {tc:>8} {nc:>8} {ratio:>8.2f}  T:{bar_t}  N:{bar_n}")
