"""Registered colexification pairs (plan/analysis-plan.md), using the family-level measure of
code/colex.py (unchanged from the earlier registered analysis of HEAVY).

For each pair (anchor, target): family-level colexification rate and its percentile among the
anchor's rates with every other concept sharing at least 30 families with it (targets of the
anchor excluded from that reference set). Anchors covered by fewer than 30 families are
reported descriptively without a percentile.
Usage: python code/pairs.py <forms.csv> <languages.csv> <out.json>
"""
import json
import sys

import numpy as np

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

PAIRS = [("DOWN", "SAD"), ("LOW", "SAD"),
         ("HARD", "ANGRY"), ("HARD", "ANGER"), ("TIGHT", "ANGRY"), ("TIGHT", "ANGER"),
         ("CROOKED", "BAD"), ("CROOKED", "EVIL"), ("CROOKED", "WRONG"),
         ("TWIST", "BAD"), ("TWIST", "EVIL"),
         ("STRAIGHT", "CORRECT (RIGHT)"), ("STRAIGHT", "RIGHT")]
POSITIVE_CONTROL = [("STRAIGHT", "CORRECT (RIGHT)")]
MIN_FAMILIES = 30


def main():
    lex, fam = colex.load(sys.argv[1], sys.argv[2])
    out = {"pairs": {}}
    anchors = sorted({a for a, _ in PAIRS})
    for anchor in anchors:
        rates = colex.rates_for(anchor, lex, fam)
        targets = {t for a, t in PAIRS if a == anchor}
        ref = [k / n for c, (k, n) in rates.items() if n >= MIN_FAMILIES and c not in targets]
        anchor_families = len({fam.get(v, "unknown:" + v) for v, c in lex.items() if anchor in c})
        for t in targets:
            k, n = rates.get(t, (0, 0))
            rate = k / n if n else float("nan")
            rec = {"families_colexifying": k, "families_with_both": n, "rate": round(rate, 4) if n else None,
                   "anchor_families": anchor_families, "reference_partners": len(ref)}
            if anchor_families >= MIN_FAMILIES and n >= MIN_FAMILIES and ref:
                pct = colex.percentile_of(rate, ref)
                rec.update({"percentile": round(pct, 4), "top_5_percent": bool(pct >= 0.95),
                            "reference_p95": round(float(np.percentile(ref, 95)), 4)})
            else:
                rec.update({"percentile": None, "top_5_percent": None, "note": "too few families for the percentile test"})
            out["pairs"][f"{anchor}|{t}"] = rec
    with open(sys.argv[3], "w") as fh:
        json.dump(out, fh, indent=2, sort_keys=True, ensure_ascii=False)
        fh.write("\n")
    print(json.dumps(out, indent=1, ensure_ascii=False))


if __name__ == "__main__":
    main()