"""Registered runner for protocol-acoustics-model.json: the four Hearing Anything Anywhere rooms.

    python run-haa.py <protocol.json> <haa-scenes.json> <holdout dir> <model dir> <out.json>

Reads the holdout responses once, scores every arm, writes the result whatever it says. Checks before
reading: the model file's sha256 equals the protocol's. The candidate is the calibrated formula carried
to every position by the learned field (fewshot.gp_predict with the lengths in kernel.json).
"""
import hashlib
import json
import math
import sys

import numpy as np
from scipy.spatial import ConvexHull

import arms
import fewshot as FS
import features as F
import tasks as T
from common import KEYS, analyse, item_score, memory_alone, with_memory
from livingworlds_map.backends import rooms as R
from simcorpus import MATERIALS as TWIN

FS_HZ = 48000
CLASS = {"plasterboard": "plasterboard", "floor": "generic", "ceiling": "generic", "absorber": "acoustic panel",
         "glass": "glass", "wood_floor": "wood", "concrete": "concrete"}


def sha(path):
    return hashlib.sha256(open(path, "rb").read()).hexdigest()


def room_of(scene):
    surf = [R.Surface(np.array(s["polygon"], float), tuple(TWIN[CLASS[s["class"]]]), s["name"]) for s in scene["surfaces"]]
    pts = np.vstack([s.vertices for s in surf])
    hull = ConvexHull(pts)
    inside = pts[hull.vertices].mean(0)
    return R.Room(surf, volume=float(hull.volume), interior=inside)


def items_of(scene, xyz, rirs):
    out = []
    for i in range(len(xyz)):
        a = analyse(np.asarray(rirs[i], float), FS_HZ)
        out.append({"src": list(scene["speaker"]), "rcv": [float(v) for v in xyz[i]], **a})
    return out


def main(protocol, scenes, holdout, modeldir, outpath):
    P = json.load(open(protocol))
    psha = sha(protocol)
    for name, want in P["model"]["sha256"].items():
        got = sha(modeldir + "/" + name)
        if got != want:
            raise SystemExit("model file %s: sha256 %s, protocol says %s" % (name, got, want))
    hyper = {k: tuple(v) for k, v in json.load(open(modeldir + "/kernel.json")).items()}
    G = json.load(open(scenes))
    result = {"protocolSha256": psha, "rooms": {}}
    pooled = {}
    for room, scene in G.items():
        cfg = scene["config"]
        xyz = np.load("%s/%s-xyzs.npy" % (holdout, cfg))
        rirs = np.load("%s/%s-RIRs.npy" % (holdout, cfg), mmap_mode="r")
        items = items_of(scene, xyz, rirs)
        task = T.build(room_of(scene), items, scene["train"], "haa-" + room)
        meas = [items[j] for j in scene["train"]]
        idx = {id(it): i for i, it in enumerate(items)}
        priors = {"formula": task["base"], "ism": arms.ism_preds(task)}
        cand = FS.gp_predict(task, task["base"], hyper)
        per = {}
        for i in scene["test"]:
            it = items[i]
            if not any(np.isfinite(it[k]).any() for k in KEYS):
                per.setdefault("did not run", []).append(i)
                continue
            s = {"memory": item_score(memory_alone(meas, it), it)}
            for pn, Pp in priors.items():
                s[pn + "+memory"] = item_score(with_memory(lambda m, Pp=Pp: Pp[idx[id(m)]], meas, it), it)
            s["formula+field"] = item_score(cand[i], it)
            for k, v in s.items():
                per.setdefault(k, []).append(v)
                pooled.setdefault(k, []).append(v)
        result["rooms"][room] = {k: (float(np.nanmean(v)) if k != "did not run" else v) for k, v in per.items()}
        result["rooms"][room]["scored"] = len(per["memory"])
        print(room, result["rooms"][room], flush=True)
    d = np.array(pooled["formula+field"]) - np.array(pooled[P["baseline"]])
    d = d[np.isfinite(d)]
    wins, losses = int((d > 0).sum()), int((d < 0).sum())
    from scipy.stats import binomtest
    p = float(binomtest(wins, wins + losses, 0.5, alternative="greater").pvalue)
    rng = np.random.default_rng(P["seed"])
    boot = [rng.choice(d, len(d)).mean() for _ in range(P["bootstrap"])]
    lo, hi = np.percentile(boot, [2.5, 97.5])
    roomwins = sum(result["rooms"][r]["formula+field"] > result["rooms"][r][P["baseline"]] for r in G)
    met = bool(p < 0.05 and lo > 0 and roomwins >= P["roomsNeeded"])
    result["pooled"] = {k: float(np.nanmean(v)) for k, v in pooled.items()}
    result["primary"] = {"meanDiff": float(d.mean()), "ci95": [float(lo), float(hi)], "wins": wins, "losses": losses,
                         "p": p, "roomsAbove": roomwins, "roomsNeeded": P["roomsNeeded"], "met": met}
    json.dump(result, open(outpath, "w"), indent=1)
    print(json.dumps(result["primary"]), flush=True)


if __name__ == "__main__":
    main(*sys.argv[1:6])