"""Reviewer spot-checks (Quiet Replication). Reuses bundle code/colex.py unchanged.
Usage (in Docker, from bundle copy root): python /ev/probe.py tmp/forms.csv data/clics4/languages.csv
"""
import csv, sys, json, math, re
from collections import defaultdict, Counter
sys.path.insert(0, "code")
import colex

forms, langs = sys.argv[1], sys.argv[2]
lex, fam = colex.load(forms, langs)
meta = {}
with open(langs, newline="", encoding="utf-8") as fh:
    for r in csv.DictReader(fh):
        meta[r["ID"]] = r
out = {}

# 1. family assignment
nofam = [v for v in lex if not meta.get(v, {}).get("Family")]
out["varieties"] = len(lex)
out["varieties_without_family"] = len(nofam)
out["distinct_families_incl_isolate_singletons"] = len({fam.get(v) for v in lex})

# 2. who colexifies each key pair
def hits(a, b):
    rows = []
    for v, c in lex.items():
        if a in c and b in c and c[a] & c[b]:
            m = meta.get(v, {})
            rows.append({"variety": m.get("Name"), "family": m.get("Family_Name") or ("isolate:" + v),
                         "contribution": m.get("Contribution_ID"), "forms": sorted(c[a] & c[b])})
    return rows
PAIRS = [("LIGHT (WEIGHT)", "EASY"), ("HEAVY", "DIFFICULT"), ("STRAIGHT", "CORRECT (RIGHT)"),
         ("STRAIGHT", "GOOD"), ("CROOKED", "WRONG"), ("CROOKED", "BAD"), ("UP", "HAPPY"),
         ("SOFT", "QUIET"), ("HARD", "ANGER")]
out["hits"] = {}
for a, b in PAIRS:
    h = hits(a, b)
    out["hits"][f"{a}|{b}"] = {"varieties": len(h), "families": dict(Counter(r["family"] for r in h)),
                              "contributions": dict(Counter(r["contribution"] for r in h)),
                              "examples": h[:60]}

# 3. coverage of the pair by contribution (which datasets carry both concepts)
def cover(a, b):
    return dict(Counter(meta.get(v, {}).get("Contribution_ID") for v, c in lex.items() if a in c and b in c).most_common(12))
out["coverage_contrib"] = {f"{a}|{b}": cover(a, b) for a, b in PAIRS[:2]}

# 4. coverage-matched reference: anchor partner rates using only varieties that carry the target
def matched(anchor, target, minfam=30):
    sub = {v: c for v, c in lex.items() if target in c}
    r = colex.rates_for(anchor, sub, fam)
    k, n = r.get(target, (0, 0))
    ref = sorted(((kk / nn, c, kk, nn) for c, (kk, nn) in r.items() if nn >= minfam and c != target), reverse=True)
    pct = colex.percentile_of(k / n, [x[0] for x in ref])
    return {"target": [k, n], "percentile_matched": round(pct, 4), "ref_n": len(ref), "top_ref": [[c, kk, nn] for _, c, kk, nn in ref[:8]]}
out["matched"] = {f"{a}|{b}": matched(a, b) for a, b in [("LIGHT (WEIGHT)", "EASY"), ("HEAVY", "DIFFICULT"),
                  ("STRAIGHT", "GOOD"), ("CROOKED", "WRONG"), ("STRAIGHT", "CORRECT (RIGHT)")]}

# 5. top partners of anchors (full data, n>=30)
def top(anchor, n=10, minfam=30):
    r = colex.rates_for(anchor, lex, fam)
    t = sorted(((k / nn, c, k, nn) for c, (k, nn) in r.items() if nn >= minfam), reverse=True)[:n]
    return [[c, k, nn] for _, c, k, nn in t]
out["top_partners"] = {a: top(a) for a in ["LIGHT (WEIGHT)", "HEAVY", "STRAIGHT", "CROOKED", "UP", "HIGH", "DOWN", "LOW", "HARD", "SOFT"]}

# 6. alternative affect/tension concepts present in the data
allg = Counter()
for v, c in lex.items():
    for g in c:
        allg[g] += 1
pat = re.compile(r"HAPP|GLAD|JOY|MERR|CHEER|SAD|SORR|GRIEF|DEPRESS|CALM|TENS|TIGHT|ANGR|ANGER|RAGE|PEACE|QUIET|GENTLE|STIFF|CORRECT|RIGHT|WRONG|GOOD|BAD|EVIL|HONEST|TRUE")
out["candidate_concepts"] = {g: allg[g] for g in sorted(allg) if pat.search(g)}
alts = {"UP": ["GLAD", "JOY", "HAPPINESS", "CHEERFUL", "MERRY"], "HIGH": ["GLAD", "JOY", "HAPPINESS", "PROUD"],
        "LOW": ["GRIEF", "SORROW", "SADNESS", "HUMBLE"], "DOWN": ["GRIEF", "SADNESS"],
        "HEAVY": ["SAD", "GRIEF", "SADNESS", "SORROW"], "LIGHT (WEIGHT)": ["GLAD", "JOY", "HAPPINESS"],
        "HARD": ["STIFF", "STRONG", "DIFFICULT", "SEVERE", "CRUEL", "STINGY"], "SOFT": ["CALM", "PEACEFUL", "WEAK", "GENTLE", "SLOW", "TENDER"],
        "STRAIGHT": ["TRUE", "HONEST", "RIGHT"], "CROOKED": ["DISHONEST", "LIE (TELL A LIE)", "FALSE"],
        "TIGHT": ["NARROW", "ANGRY"]}
alt_out = {}
for a, ts in alts.items():
    r = colex.rates_for(a, lex, fam)
    alt_out[a] = {t: list(r.get(t, (0, 0))) for t in ts if t in allg}
out["alternative_targets"] = alt_out

# 7. Fisher exact (two-sided) on family counts for asymmetry contrasts
def lchoose(n, k): return math.lgamma(n + 1) - math.lgamma(k + 1) - math.lgamma(n - k + 1)
def fisher(a, n1, b, n2):
    K, N = a + b, n1 + n2
    def p(x): return math.exp(lchoose(n1, x) + lchoose(n2, K - x) - lchoose(N, K))
    p0 = p(a)
    return sum(p(x) for x in range(max(0, K - n2), min(K, n1) + 1) if p(x) <= p0 * (1 + 1e-9))
out["fisher_two_sided"] = {
    "straight-correct 38/155 vs crooked-wrong 4/113": round(fisher(38, 155, 4, 113), 6),
    "straight-good 7/200 vs crooked-bad 1/171": round(fisher(7, 200, 1, 171), 4),
    "light-easy 11/50 vs heavy-difficult 7/57": round(fisher(11, 50, 7, 57), 4),
}
# 8. upper 95% one-sided bound for zero / one hit (exact binomial, family level)
def ub(k, n, a=0.05):
    lo, hi = k / n, 1.0
    for _ in range(60):
        m = (lo + hi) / 2
        cdf = sum(math.comb(n, i) * m**i * (1 - m)**(n - i) for i in range(k + 1))
        lo, hi = (m, hi) if cdf > a else (lo, m)
    return round(hi, 4)
out["upper95"] = {"up-happy 1/66": ub(1, 66), "high-happy 0/145": ub(0, 145), "down-sad 0/21": ub(0, 21),
                  "low-sad 0/108": ub(0, 108), "light-happy 0/146": ub(0, 146), "hard-angry 0/168": ub(0, 168),
                  "soft-quiet 2/154": ub(2, 154)}
# 9. baseline: share of eligible reference partners with rate 0 for each null anchor
def zeroshare(anchor):
    r = colex.rates_for(anchor, lex, fam)
    rates = [k / n for c, (k, n) in r.items() if n >= 30]
    return {"partners": len(rates), "share_zero": round(sum(x == 0 for x in rates) / len(rates), 4)}
out["baseline_zero_share"] = {a: zeroshare(a) for a in ["UP", "HIGH", "LOW", "LIGHT (WEIGHT)", "HARD", "SOFT"]}
json.dump(out, open("/ev/probe_out.json", "w"), indent=1, ensure_ascii=False)
print(json.dumps({k: v for k, v in out.items() if k != "hits"}, indent=1, ensure_ascii=False))
for k, v in out["hits"].items():
    print(k, v["varieties"], v["families"], v["contributions"])