#!/usr/bin/env python3
"""build_hivemind.py — answer distributions for "The Hive Mind Test" (P17 front-piece).

For each synesthetic probe (what colour is October? what temperature is anger?), parse the
model's answer into a canonical bucket and tally the distribution across all families. The
quiz grades a human's pick against these real distributions: pick the dominant answer and
you match the hive mind; pick a rare one and you're an outlier. Writes hivemind.json.

Pure stdlib. Re-run:  python3 research/the-hive-mind-test/build_hivemind.py
Sanity-asserts the dominant answers against the committed finding (results/synesthesia.md):
October→orange, letter-A→red, anger→hot.
"""
import json, glob, collections, re
from pathlib import Path

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]
RAW = sorted((ROOT / "research/the-map/data/raw").glob("*.jsonl"))

# canonical answer vocabularies; first synonym found in the output (lowercased) wins.
COLOUR = {
    "red": ["red", "crimson", "scarlet", "ruby", "maroon"],
    "orange": ["orange", "amber", "rust", "burnt"],
    "yellow": ["yellow", "gold", "golden"],
    "green": ["green", "emerald", "olive", "lime"],
    "blue": ["blue", "azure", "navy", "cyan", "indigo", "teal", "cobalt"],
    "purple": ["purple", "violet", "lavender", "magenta", "plum", "lilac"],
    "grey": ["grey", "gray", "silver", "slate"],
    "brown": ["brown", "tan", "beige", "sepia", "mahogany", "bronze"],
    "black": ["black", "obsidian", "charcoal"],
    "white": ["white", "cream", "ivory", "pearl"],
    "pink": ["pink", "rose", "salmon", "blush"],
}
TEMP = {"hot": ["hot", "heat", "burning", "fire", "scorching", "boiling", "warm", "fever", "blazing", "molten"],
        "cold": ["cold", "cool", "freezing", "icy", "ice", "chilly", "frigid", "frozen"]}
SHAPE = {"round": ["round", "circle", "circular", "sphere", "spiral", "curve", "oval", "bubble"],
         "angular": ["angular", "spiky", "jagged", "sharp", "star", "triangle", "zigzag", "burst"],
         "square": ["square", "rectangle", "cube", "block"],
         "wavy": ["wavy", "undulating", "wave", "ripple", "squiggle"]}

QUESTIONS = [
    ("syn-october-colour",   "October",        "colour", COLOUR),
    ("syn-letterA-colour",   "the letter A",   "colour", COLOUR),
    ("syn-loneliness-colour","loneliness",     "colour", COLOUR),
    ("syn-cello-colour",     "a cello",        "colour", COLOUR),
    ("syn-number7-colour",   "the number 7",   "colour", COLOUR),
    ("syn-justice-colour",   "justice",        "colour", COLOUR),
    ("syn-wednesday-colour", "Wednesday",      "colour", COLOUR),
    ("syn-anger-temperature","anger",          "temperature", TEMP),
    ("syn-joy-shape",        "joy",            "shape", SHAPE),
]

def classify(text, vocab):
    low = text.lower()
    best, bestpos = None, len(low) + 1
    for canon, syns in vocab.items():
        for s in syns:
            m = re.search(r"\b" + re.escape(s) + r"\b", low)
            if m and m.start() < bestpos:
                best, bestpos = canon, m.start()
    return best

by_probe = collections.defaultdict(lambda: {"dist": collections.Counter(), "fams": set(), "n": 0, "unparsed": 0})
want = {q[0] for q in QUESTIONS}
for f in RAW:
    for line in open(f, encoding="utf-8"):
        line = line.strip()
        if not line:
            continue
        r = json.loads(line)
        pid = r.get("probe_id")
        if pid not in want:
            continue
        vocab = next(q[3] for q in QUESTIONS if q[0] == pid)
        rec = by_probe[pid]
        out = (r.get("output_text") or "").strip()
        if not out:
            continue
        rec["n"] += 1; rec["fams"].add(r.get("family"))
        a = classify(out, vocab)
        if a: rec["dist"][a] += 1
        else: rec["unparsed"] += 1

questions = []
for pid, label, atype, vocab in QUESTIONS:
    rec = by_probe[pid]
    total = sum(rec["dist"].values())
    dist = {a: c for a, c in rec["dist"].most_common()}
    top, topc = rec["dist"].most_common(1)[0]
    questions.append({
        "id": pid, "label": label, "type": atype,
        "options": list(vocab.keys()),
        "dist": dist, "total": total, "families": len(rec["fams"]),
        "top": top, "topPct": round(100 * topc / total),
        "parsedPct": round(100 * total / (total + rec["unparsed"])) if (total + rec["unparsed"]) else 0,
    })
    print(f"{label:14}({atype:11}) n={total:4} parsed={questions[-1]['parsedPct']}% | "
          f"top={top} {questions[-1]['topPct']}% | " +
          " ".join(f"{a}:{round(100*c/total)}" for a, c in rec["dist"].most_common(5)))

out = {"source": "The Map's synesthetic probes", "families": max(q["families"] for q in questions),
       "note": "Each question's distribution is real model answers parsed into canonical buckets; "
               "the quiz grades a human pick against it. Reproduce: build_hivemind.py.",
       "questions": questions}
(HERE / "hivemind.json").write_text(json.dumps(out, ensure_ascii=False, indent=1), encoding="utf-8")
print(f"\nwrote {HERE/'hivemind.json'}")

# --- sanity vs the committed finding (results/synesthesia.md) ---
byid = {q["id"]: q for q in questions}
assert byid["syn-october-colour"]["top"] == "orange", "October should be orange"
assert byid["syn-letterA-colour"]["top"] == "red", "letter A should be red"
assert byid["syn-anger-temperature"]["top"] == "hot", "anger should be hot"
print("OK — dominant answers match the committed synesthesia finding (October→orange, A→red, anger→hot)")