"""Correlation analysis: how do individual features relate to each other?

Computes all feature scores for each pair, then shows:
- Correlation matrix
- Which features provide independent signal
- Where features disagree (one high, another low)
"""

import csv
import unicodedata
import re
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"
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 bigrams(s): return set(s[i:i+2] for i in range(len(s) - 1))
def dice(a, b):
    if not a and not b: return 1.0
    if not a or not b: return 0.0
    return 2.0 * len(a & b) / (len(a) + len(b))

# Shape class lookup for weighted lev
SHAPE_CLASS = {}
classes = [("V","il1I|!tfj"),("O","oO0cCeQa"),("M","mwMW"),("N","nuvNUV"),
           ("B","bhk6"),("P","pqgy9"),("D","dDPRB"),("X","xXzZK"),("S","sS5$")]
for cls, chars in classes:
    for ch in chars: SHAPE_CLASS[ch.lower()] = cls

def weighted_lev(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):
            cost = 0.0 if s1[i-1]==s2[j-1] else (0.2 if SHAPE_CLASS.get(s1[i-1])==SHAPE_CLASS.get(s2[j-1]) and SHAPE_CLASS.get(s1[i-1]) else 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 compute_features(query, candidate):
    q, c = norm(query), norm(candidate)
    if not q or not c:
        return {k: 0.0 for k in ["jw","rev_jw","lev","shape_lev","shape_dice","weighted_lev","len_ratio"]}

    sq, sc = shape_compress(q), shape_compress(c)
    wl_dist = weighted_lev(q, c)
    max_len = max(len(q), len(c))

    return {
        "jw": JaroWinkler.similarity(q, c),
        "rev_jw": JaroWinkler.similarity(q[::-1], c[::-1]),
        "lev": fuzz.ratio(q, c) / 100.0,
        "shape_lev": fuzz.ratio(sq, sc) / 100.0,
        "shape_dice": dice(bigrams(sq), bigrams(sc)) if len(sq)>=2 and len(sc)>=2 else (1.0 if sq==sc else 0.0),
        "weighted_lev": max(0.0, 1.0 - wl_dist / max_len) if max_len > 0 else 0.0,
        "len_ratio": length_ratio(q, c),
    }


def pearson(x, y):
    n = len(x)
    mx, my = sum(x)/n, sum(y)/n
    sx = (sum((xi-mx)**2 for xi in x) / n) ** 0.5
    sy = (sum((yi-my)**2 for yi in y) / n) ** 0.5
    if sx == 0 or sy == 0:
        return 0.0
    return sum((xi-mx)*(yi-my) for xi, yi in zip(x, y)) / (n * sx * sy)


if __name__ == "__main__":
    pairs = []
    with open(CSV_PATH) as f:
        for row in csv.DictReader(f):
            pairs.append((row["ocr_error"], row["correction"]))

    feature_names = ["jw", "rev_jw", "lev", "shape_lev", "shape_dice", "weighted_lev", "len_ratio"]
    all_features = {k: [] for k in feature_names}

    for ocr_error, correction in pairs:
        feats = compute_features(correction, ocr_error)
        for k in feature_names:
            all_features[k].append(feats[k])

    # Correlation matrix
    print("=" * 80)
    print("  CORRELATION MATRIX (Pearson)")
    print("=" * 80)
    header = f"{'':>14}" + "".join(f"{n:>14}" for n in feature_names)
    print(header)
    for a in feature_names:
        row = f"{a:>14}"
        for b in feature_names:
            r = pearson(all_features[a], all_features[b])
            row += f"{r:>14.3f}"
        print(row)

    # Mean scores
    print(f"\n{'':>14}" + "".join(f"{n:>14}" for n in feature_names))
    print(f"{'mean':>14}" + "".join(f"{sum(all_features[n])/len(all_features[n]):>14.3f}" for n in feature_names))

    # Find pairs where features DISAGREE most (one high, another low)
    print("\n" + "=" * 80)
    print("  DISAGREEMENTS: JW high but Shape low (JW - Shape > 0.15)")
    print("=" * 80)
    print(f"  {'OCR Error':<20} {'Correction':<20} {'JW':>6} {'RevJW':>6} {'ShpLv':>6} {'ShpDc':>6} {'WtLev':>6} {'LenR':>6}")
    print(f"  {'-'*20} {'-'*20} {'-'*6} {'-'*6} {'-'*6} {'-'*6} {'-'*6} {'-'*6}")
    cases = []
    for i, (ocr, corr) in enumerate(pairs):
        jw = all_features["jw"][i]
        sl = all_features["shape_lev"][i]
        if jw - sl > 0.15:
            cases.append((jw - sl, ocr, corr, i))
    cases.sort(reverse=True)
    for _, ocr, corr, i in cases[:15]:
        print(f"  {ocr:<20} {corr:<20} {all_features['jw'][i]:>6.3f} {all_features['rev_jw'][i]:>6.3f} {all_features['shape_lev'][i]:>6.3f} {all_features['shape_dice'][i]:>6.3f} {all_features['weighted_lev'][i]:>6.3f} {all_features['len_ratio'][i]:>6.3f}")

    print("\n" + "=" * 80)
    print("  DISAGREEMENTS: Shape high but JW low (Shape - JW > 0.15)")
    print("=" * 80)
    print(f"  {'OCR Error':<20} {'Correction':<20} {'JW':>6} {'RevJW':>6} {'ShpLv':>6} {'ShpDc':>6} {'WtLev':>6} {'LenR':>6}")
    print(f"  {'-'*20} {'-'*20} {'-'*6} {'-'*6} {'-'*6} {'-'*6} {'-'*6} {'-'*6}")
    cases2 = []
    for i, (ocr, corr) in enumerate(pairs):
        jw = all_features["jw"][i]
        sl = all_features["shape_lev"][i]
        if sl - jw > 0.15:
            cases2.append((sl - jw, ocr, corr, i))
    cases2.sort(reverse=True)
    for _, ocr, corr, i in cases2[:15]:
        print(f"  {ocr:<20} {corr:<20} {all_features['jw'][i]:>6.3f} {all_features['rev_jw'][i]:>6.3f} {all_features['shape_lev'][i]:>6.3f} {all_features['shape_dice'][i]:>6.3f} {all_features['weighted_lev'][i]:>6.3f} {all_features['len_ratio'][i]:>6.3f}")

    print("\n" + "=" * 80)
    print("  DISAGREEMENTS: RevJW rescues (RevJW - JW > 0.05)")
    print("=" * 80)
    print(f"  {'OCR Error':<20} {'Correction':<20} {'JW':>6} {'RevJW':>6} {'ShpLv':>6} {'Diff':>6}")
    print(f"  {'-'*20} {'-'*20} {'-'*6} {'-'*6} {'-'*6} {'-'*6}")
    cases3 = []
    for i, (ocr, corr) in enumerate(pairs):
        jw = all_features["jw"][i]
        rjw = all_features["rev_jw"][i]
        if rjw - jw > 0.05:
            cases3.append((rjw - jw, ocr, corr, i))
    cases3.sort(reverse=True)
    for _, ocr, corr, i in cases3[:15]:
        diff = all_features["rev_jw"][i] - all_features["jw"][i]
        print(f"  {ocr:<20} {corr:<20} {all_features['jw'][i]:>6.3f} {all_features['rev_jw'][i]:>6.3f} {all_features['shape_lev'][i]:>6.3f} {diff:>+6.3f}")
