"""Registered screen (plan/analysis-plan.md): do words for mechanical stress, motion and path
double as words for emotion, cognition or effort across language families?

Measure (code/colex.py, unchanged from earlier registered analyses): family-level colexification
rate of an anchor with a target, and its percentile among the anchor's rates with every
non-target concept sharing at least MIN_FAMILIES families with it. A pair is a hit if its
percentile is at least 0.95 and at least MIN_HIT_FAMILIES families colexify it.

Class test: the number of hits (bodily targets excluded) among a class's anchors, against the
number for random sets of the same size drawn from a fixed random pool of other concepts.

Usage: python code/mechanics.py <forms.csv> <languages.csv> <out.json>
"""
import csv
import json
import sys
from collections import defaultdict

import numpy as np

sys.path.insert(0, "code")
import colex  # noqa: E402

CLASSES = {
    "load": ["HEAVY", "LIGHT (WEIGHT)", "CARRY", "LIFT", "WEIGH"],
    "normal_stress": ["PRESS", "SQUEEZE", "CRUSH", "PUSH", "PULL", "STRETCH", "SWELL", "BREAK (BREAKING)"],
    "shear": ["RUB", "SCRAPE", "SCRATCH", "GRIND", "TEAR (SHRED)", "SPLIT", "TWIST", "BEND", "ROUGH", "SMOOTH", "CUT"],
    "flow_momentum": ["FLOW", "FAST", "SLOW", "HURRY", "MOVE", "FALL", "SINK (DESCEND)", "RISE (MOVE UPWARDS)",
                      "FLOAT", "SHAKE", "SPIN", "BOIL (OF LIQUID)"],
    "path": ["PATH", "ROAD", "FOLLOW", "TURN AROUND", "COME BACK", "TRAP (PITFALL)"],
}
TARGETS = {
    "affect": ["SAD", "GRIEF", "ANGRY", "ANGER", "FEAR (BE AFRAID)", "FEAR (FRIGHT)", "ANXIETY", "SHAME", "HATE",
               "ENVY", "PITY", "REGRET", "GUILTY", "HAPPY", "LOVE", "PEACE", "HOPE", "PROUD", "SURPRISED", "BRAVE"],
    "cognition": ["FORGET", "REMEMBER", "UNDERSTAND", "THINK (REFLECT)", "THINK (BELIEVE)", "BELIEVE", "DOUBT", "WANT"],
    "effort": ["DIFFICULT", "EASY"],
    "bodily": ["PAIN", "PAINFUL", "TIRED", "LAZY"],
}
ALL_TARGETS = [t for ts in TARGETS.values() for t in ts]
COUNTED = [t for k, ts in TARGETS.items() if k != "bodily" for t in ts]
MIN_FAMILIES = 30
MIN_HIT_FAMILIES = 3
POOL_MIN_FAMILIES = 60
POOL_SIZE = 200
N_NULL = 10000
SEED = 20261009


def screen(anchor, lex, fam):
    rates = colex.rates_for(anchor, lex, fam)
    ref = [k / n for c, (k, n) in rates.items() if n >= MIN_FAMILIES and c not in ALL_TARGETS]
    out = {}
    for t in ALL_TARGETS:
        k, n = rates.get(t, (0, 0))
        if n < MIN_FAMILIES or not ref:
            continue
        pct = colex.percentile_of(k / n, ref)
        out[t] = {"families_colexifying": k, "families_with_both": n, "rate": round(k / n, 4),
                  "percentile": round(pct, 4), "hit": bool(pct >= 0.95 and k >= MIN_HIT_FAMILIES)}
    return out


def hits(res):
    return sum(1 for t, r in res.items() if r["hit"] and t in COUNTED)


def witnesses(anchor, target, lex, fam, area):
    """Varieties colexifying a hit pair, with family and macroarea (descriptive)."""
    out = []
    for var, cs in lex.items():
        if anchor in cs and target in cs:
            shared = cs[anchor] & cs[target]
            if shared:
                out.append({"variety": var, "family": fam.get(var), "macroarea": area.get(var), "forms": sorted(shared)})
    return out


def main():
    forms, langs, out_path = sys.argv[1], sys.argv[2], sys.argv[3]
    lex, fam = colex.load(forms, langs)
    area, names = {}, {}
    with open(langs, newline="", encoding="utf-8") as fh:
        for r in csv.DictReader(fh):
            area[r["ID"]] = r["Macroarea"]
            names[r["ID"]] = r["Name"]
    cover = defaultdict(set)
    for v, cs in lex.items():
        for c in cs:
            cover[c].add(fam.get(v, "unknown:" + v))
    anchors = [a for cl in CLASSES.values() for a in cl]
    rng = np.random.default_rng(SEED)
    candidates = sorted(c for c, s in cover.items() if len(s) >= POOL_MIN_FAMILIES and c not in anchors and c not in ALL_TARGETS)
    pool = list(rng.choice(candidates, POOL_SIZE, replace=False))

    res = {"classes": {}, "pairs": {}, "hits": []}
    anchor_hits = {}
    for a in anchors:
        s = screen(a, lex, fam)
        res["pairs"][a] = s
        anchor_hits[a] = hits(s)
    pool_hits = np.array([hits(screen(p, lex, fam)) for p in pool])
    res["pool"] = {"size": POOL_SIZE, "mean_hits_per_concept": round(float(pool_hits.mean()), 4),
                   "share_with_any_hit": round(float((pool_hits > 0).mean()), 4)}
    pvals = {}
    for cl, members in CLASSES.items():
        obs = int(sum(anchor_hits[a] for a in members))
        null = np.array([pool_hits[rng.choice(POOL_SIZE, len(members), replace=False)].sum() for _ in range(N_NULL)])
        p = float(((null >= obs).sum() + 1) / (N_NULL + 1))
        pvals[cl] = p
        res["classes"][cl] = {"anchors": len(members), "hits": obs, "null_mean": round(float(null.mean()), 3),
                              "null_p95": float(np.percentile(null, 95)), "p_one_sided": round(p, 4),
                              "hits_per_anchor": round(obs / len(members), 3)}
    # Holm over the four non-benchmark classes
    tested = sorted((p, cl) for cl, p in pvals.items() if cl != "load")
    m = len(tested)
    running = 0.0
    for i, (p, cl) in enumerate(tested):
        running = max(running, min(1.0, (m - i) * p))
        res["classes"][cl]["p_holm"] = round(running, 4)
    for cl, members in CLASSES.items():
        for a in members:
            for t, r in res["pairs"][a].items():
                if r["hit"]:
                    w = witnesses(a, t, lex, fam, area)
                    tcls = next(k for k, ts in TARGETS.items() if t in ts)
                    res["hits"].append({"class": cl, "anchor": a, "target": t, "target_class": tcls,
                                        "families": r["families_colexifying"], "of": r["families_with_both"],
                                        "percentile": r["percentile"],
                                        "macroareas": sorted({x["macroarea"] for x in w if x["macroarea"]}),
                                        "examples": [f"{names.get(x['variety'], x['variety'])}: {', '.join(x['forms'])}" for x in w[:8]]})
    with open(out_path, "w") as fh:
        json.dump(res, fh, indent=2, sort_keys=True, ensure_ascii=False)
        fh.write("\n")
    print(json.dumps({"classes": res["classes"], "pool": res["pool"]}, indent=1))
    for h in res["hits"]:
        print(h["class"], h["anchor"], "|", h["target"], f"{h['families']}/{h['of']}", h["percentile"], h["macroareas"])


if __name__ == "__main__":
    main()