"""Registered scorer (protocol-peds.json). For each scene, four crowds against the real one:
jupedsim, socialforce, orca (from the render session's pedreplay.py) and straight (peds_straight.py).
Four measures, lower is closer: 1-Wasserstein distance of step speeds, of nearest-neighbour distances
and of path straightness (peds_realism.py), and |share of frames with a contact - the real share|.
Each measure ranks the four crowds (ties share the mean rank); a crowd's score in a scene is its
rank sum. Usage: peds_score.py SIMDIR SCENE...   (SIMDIR holds <scene>-<model>.txt)"""
import json, os, sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from peds_realism import load, stats, w1

MODELS = ["jupedsim", "socialforce", "orca", "straight"]
simdir, scenes = sys.argv[1], sys.argv[2:]
here = os.path.dirname(os.path.abspath(__file__))


def ranks(vals):
    order = sorted(range(len(vals)), key=lambda i: vals[i]); r = [0.0] * len(vals); i = 0
    while i < len(order):
        j = i
        while j + 1 < len(order) and vals[order[j + 1]] == vals[order[i]]: j += 1
        for k in range(i, j + 1): r[order[k]] = (i + j) / 2 + 1
        i = j + 1
    return r


out, beats_sf_and_straight, beats_orca = {}, 0, 0
for sc in scenes:
    real = stats(load(os.path.join(here, f"{sc}.txt")))
    m = {}
    for mod in MODELS:
        s = stats(load(os.path.join(simdir, f"{sc}-{mod}.txt")))
        m[mod] = {"speed": w1(real["speed"], s["speed"]), "nearest": w1(real["nearest"], s["nearest"]), "straight": w1(real["straight"], s["straight"]),
                  "contact": abs(s["contact"] - real["contact"])}
    score = {mod: 0.0 for mod in MODELS}
    for key in ("speed", "nearest", "straight", "contact"):
        for mod, r in zip(MODELS, ranks([round(m[mod][key], 6) for mod in MODELS])): score[mod] += r
    win = score["jupedsim"] < score["socialforce"] and score["jupedsim"] < score["straight"]
    beats_sf_and_straight += win; beats_orca += score["jupedsim"] < score["orca"]
    out[sc] = {"measures": {k: {q: round(v, 4) for q, v in d.items()} for k, d in m.items()}, "rankSum": score, "jupedsimBeatsSocialforceAndStraight": win,
               "contactReal": round(real["contact"], 4)}
need = 5 if len(scenes) == 6 else len(scenes) - 1
out["summary"] = {"scenes": len(scenes), "jupedsimBeatsSocialforceAndStraight": beats_sf_and_straight, "needed": need, "jupedsimBeatsOrca": beats_orca,
                  "primary_met": bool(beats_sf_and_straight >= need)}
print(json.dumps(out, indent=1))