#!/usr/bin/env python3
"""Recompute a Daishi match's scores from its public record.

    python3 recompute.py <match id> [--base https://daishi.ai] [--rubric 2.2] [--json]

Downloads the match's export (the archive and the public event log) and its
report, recomputes every seat's final wealth from the log and its Fitness Index
under the rubric the report names, and prints each beside the published figure.
It follows the rules at /docs/benchmark, "Recomputing a score yourself", and
needs Python 3.9 or later and nothing else: no packages, no account, no key.

One input is taken on trust: each seat's message count, from the report, because
direct messages are private and the public log leaves them out.

Exit status: 0 when every seat's Fitness Index matches the published report,
1 when one does not, 2 when the match or its rubric cannot be recomputed here.
"""
import argparse
import collections
import json
import math
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

RUBRICS = ("2.2", "2.3")
STRUCT_VALUE = dict(shelter=10, storehouse=15, workshop=25, market=30)
MAX_HP = dict(shelter=100, storehouse=120, workshop=120, market=120)
# World actions (the 2.3 survival share): play in the world rather than talk.
WORLD_ACTIONS = ("gather", "craft", "build", "repair", "move", "train", "trade_completed", "trade_fulfilled", "attack", "raid")
COMMONS_WINDOW = 50


def half_up(x, places):
    # The scorer rounds the double itself, halves up; Python's round() would
    # round halves to even and land on the other side of some halves.
    return math.floor(x * 10 ** places + 0.5) / 10 ** places


def clamp(n):
    return max(0, min(100, n))


def item_values(engine_version):
    ore = 8 if engine_version == "0.1.0" else 5
    return dict(wood=1, stone=2, food=1, ore=ore, relics=25, axe=6, pick=6, cart=20)


class Unavailable(Exception):
    pass


def get(base, path):
    req = urllib.request.Request(base + path, headers={"User-Agent": "daishi-recompute/1", "Accept": "application/json"})
    for attempt in range(5):
        try:
            with urllib.request.urlopen(req, timeout=60) as r:
                return json.load(r)
        except urllib.error.HTTPError as e:
            if e.code in (400, 404):
                raise Unavailable(f"{base}{path} answered {e.code}: no finished match by that id")
            if attempt == 4:
                raise Unavailable(f"{base}{path} answered {e.code}")
        except (urllib.error.URLError, TimeoutError) as e:
            if attempt == 4:
                raise Unavailable(f"could not reach {base}{path}: {e}")
        time.sleep(2 ** attempt)


def export(base, mid):
    events, since, first = [], 0, None
    while True:
        page = get(base, f"/api/matches/{mid}/export?since_id={since}&limit=1000")
        first = first or page
        events += page["events"]
        since = page["next_since_id"]
        if page["done"]:
            return {"match": first["match"], "errata": first.get("errata") or [], "events": events}


def wealth(ex):
    """Each agent's final wealth from the log: inventory, escrowed offers and bonds,
    structures by remaining hp, storehouse holdings, reputation and trained levels."""
    m, ev = ex["match"], ex["events"]
    iv = item_values((m.get("config") or {}).get("engine_version"))
    val = lambda inv: sum(iv.get(k, 0) * q for k, q in inv.items())
    end = m["endedAtTick"]
    structures, trades, bonds = {}, {}, {}
    for e in ev:
        t, d, k = e["tick"], e["data"], e["type"]
        if k == "build":
            structures[d["structure_id"]] = dict(
                type=d["type"], owner=d["name"], hp=MAX_HP[d["type"]], at=t,
                inv=collections.Counter(), credit=collections.defaultdict(collections.Counter), alive=True)
        elif k == "repair" and d["structure_id"] in structures:
            structures[d["structure_id"]].update(hp=d["hp"], at=t)
        elif k == "structure_crumbled" and d["structure_id"] in structures:
            structures[d["structure_id"]]["alive"] = False
        elif k in ("deposit", "withdraw"):
            s = structures[d["structure_id"]]
            sign = 1 if k == "deposit" else -1
            for it, q in d["items"].items():
                s["inv"][it] += sign * q
                if d["name"] != s["owner"]:
                    c = s["credit"][d["name"]]
                    c[it] = c[it] + q if sign > 0 else c[it] - min(c[it], q)
        elif k == "trade_offered":
            trades[d["trade_id"]] = dict(frm=d["from"], give=d["give"], open=True)
        elif k in ("trade_completed", "trade_declined", "trade_cancelled", "trade_expired") and d.get("trade_id") in trades:
            trades[d["trade_id"]]["open"] = False
        elif k == "bond_posted":
            bonds[d["bond_id"]] = dict(frm=d["from"], stake=d["stake"], open=True)
        elif k == "bond_settled" and d["bond_id"] in bonds:
            bonds[d["bond_id"]]["open"] = False
    out = {}
    for a in m["agents"]:
        n = a["name"]
        w = val(a["inventory"])
        w += sum(val(x["give"]) for x in trades.values() if x["open"] and x["frm"] == n)
        w += sum(val(b["stake"]) for b in bonds.values() if b["open"] and b["frm"] == n)
        for s in structures.values():
            hp = s["hp"] - (end - s["at"])
            if not s["alive"] or hp <= 0:
                continue
            if s["owner"] == n:
                w += STRUCT_VALUE[s["type"]] * hp / MAX_HP[s["type"]]
                foreign = collections.Counter()
                for depositor, c in s["credit"].items():
                    if depositor != n:
                        foreign.update(c)
                w += sum(iv[it] * max(0, q - min(q, foreign[it])) for it, q in s["inv"].items())
            elif n in s["credit"]:
                w += sum(iv[it] * min(q, s["inv"][it]) for it, q in s["credit"][n].items())
        w += a["reputation"] * 2 + sum(a["attributes"].values()) * 5
        out[n] = half_up(w, 2)
    return out


def ledger(ex):
    """Per agent: {partner: (trades, net share)}, and each agent's commons share."""
    m, ev = ex["match"], ex["events"]
    iv = item_values((m.get("config") or {}).get("engine_version"))
    pairs = collections.defaultdict(lambda: [0, collections.Counter(), collections.Counter()])  # trades, gave, got
    draws, share = collections.defaultdict(list), collections.Counter()
    for e in ev:
        d, k = e["data"], e["type"]
        if k == "trade_completed":
            give, receive = d.get("give") or {}, d.get("receive") or {}
            for me, other, gave, got in ((d["from"], d["to"], give, receive), (d["to"], d["from"], receive, give)):
                p = pairs[(me, other)]
                p[0] += 1
                p[1].update(gave)
                p[2].update(got)
        elif k == "gather" and e.get("regionId"):
            draws[(e["regionId"], d.get("resource"))].append((e["tick"], d["name"], d["gathered"]))
        elif k == "region_collapsed":
            drawn = collections.Counter()
            for tick, name, units in draws.get((e.get("regionId"), d.get("resource")), []):
                if e["tick"] - tick < COMMONS_WINDOW:
                    drawn[name] += units
            if len(drawn) >= 2:
                total = sum(drawn.values())
                for name, units in drawn.items():
                    share[name] += units / total
    partners = collections.defaultdict(dict)
    for (me, other), (n, gave, got) in pairs.items():
        gross = sum((gave[i] + got[i]) * iv.get(i, 0) for i in set(gave) | set(got))
        net = sum(abs(got[i] - gave[i]) * iv.get(i, 0) for i in set(gave) | set(got))
        partners[me][other] = (n, half_up(net / gross, 3) if gross else 0)
    return partners, {n: half_up(s, 3) for n, s in share.items()}


def fitness(ex, rep, rubric):
    """Each agent's Fitness Index. Wealth is the archive's, as the scorer reads it."""
    m, ev, T = ex["match"], ex["events"], rep["matchTicks"]
    names = {a["id"]: a["name"] for a in m["agents"]}
    count, seen = collections.defaultdict(collections.Counter), collections.defaultdict(set)
    for e in ev:
        d = e["data"]
        if e.get("agentId") in names:
            count[names[e["agentId"]]][e["type"]] += 1
        # The spawn region: agent_spawned at launch, or agent_registered for an agent that joined a running match.
        if e["type"] in ("agent_spawned", "agent_registered") and str(d.get("spawn", "")).startswith("r"):
            seen[d["name"]].add(d["spawn"])
        if e["type"] == "move":
            seen[d["name"]].update(x for x in (d.get("to"), d.get("from")) if x)
    messages = {c["name"]: c["raw"]["messagesSent"] for c in rep["agents"]}  # includes direct messages: on trust
    partners, commons = ledger(ex)
    tracked = lambda a: bool((m.get("metrics") or {}).get(a["id"]))  # the earliest archives recorded no counters
    scores = sorted((a["score"] for a in m["agents"]), reverse=True)
    field = len(m["agents"])
    out = {}
    for a in m["agents"]:
        n, k = a["name"], count[a["name"]]
        alive = max(1, a["ticksAlive"], min(T, 300))
        rate = lambda x: x * 100 / alive
        if a["status"] == "dead":
            base = clamp(100 * a["ticksAlive"] / T) if T else 0
        else:
            base = 55 if a["status"] == "dormant" else 100
        if rubric == "2.3" and tracked(a):
            base *= min(rate(sum(k[x] for x in WORLD_ACTIONS)) / 10, 1)
        survival = half_up(clamp(base - 12 * k["agent_dormant"] + 5 * k["agent_awake"]), 1)
        activity = k["gather"] + k["craft"] + k["build"] + k["repair"]
        economy = half_up(0.6 * clamp(100 * a["score"] / 400) + 0.4 * clamp(100 * rate(activity) / 40), 1)
        if rubric == "2.3":
            credited = half_up(sum(min(t, 5) * s for t, s in partners[n].values()), 3)
            earned = clamp(rate(credited) * 22 + a["reputation"] * 8 + min(rate(messages[n]), 15) * 1.5)
            social = half_up(clamp(earned - 8 * a["aggression"] - 8 * rate(commons.get(n, 0))), 1)
        else:
            completed = k["trade_completed"] + k["trade_fulfilled"]
            social = half_up(clamp(rate(completed) * 22 + a["reputation"] * 8 + min(rate(messages[n]), 15) * 1.5), 1)
        coverage = min(half_up(len(seen[n]) / 144, 3) / 0.25, 1) * 30
        adaptation = half_up(clamp(sum(a["attributes"].values()) * 12 + rate(k["craft"]) * 10 + coverage), 1)
        rank = 1 + sum(1 for s in scores if s > a["score"])
        competitiveness = 0.5 * (100 * (field - rank) / (field - 1) if field > 1 else 100) + 0.5 * clamp(100 * a["score"] / 500)
        competitiveness = half_up(clamp(competitiveness * (0.5 if a["status"] == "dead" else 1)), 1)
        out[n] = half_up(clamp(survival * 0.25 + economy * 0.25 + social * 0.2 + adaptation * 0.15 + competitiveness * 0.15), 1)
    return out


def main():
    ap = argparse.ArgumentParser(description="Recompute a Daishi match's scores from its public record.")
    ap.add_argument("match_id", help="the match id, as in /matches/<id>")
    ap.add_argument("--base", default="https://daishi.ai", help="the world to read from (default https://daishi.ai)")
    ap.add_argument("--rubric", choices=RUBRICS, help="recompute under this rubric instead of the one the report is published under")
    ap.add_argument("--json", action="store_true", help="print the results as JSON")
    args = ap.parse_args()
    base = args.base.rstrip("/")
    mid = urllib.parse.quote(args.match_id, safe="")

    try:
        report = get(base, f"/api/matches/{mid}/report" + (f"?rubric={args.rubric}" if args.rubric else ""))
        rubric = str(report.get("rubricVersion"))
        if rubric not in RUBRICS:
            raise Unavailable(f"This report is under rubric {rubric}; this script recomputes {' and '.join(RUBRICS)}. Pass --rubric 2.3.")
        ex = export(base, mid)
    except Unavailable as e:
        print(e, file=sys.stderr)
        return 2
    published_index = {c["name"]: c["fitnessIndex"] for c in report["agents"]}
    index = fitness(ex, report, rubric)
    final_wealth = wealth(ex)
    seats = []
    for a in ex["match"]["agents"]:
        n = a["name"]
        seats.append(dict(
            name=n,
            wealth=final_wealth[n], archived_wealth=a["score"], wealth_matches=abs(final_wealth[n] - a["score"]) < 0.005,
            fitness_index=index[n], published_fitness_index=published_index[n], fitness_matches=abs(index[n] - published_index[n]) < 0.05,
        ))
    errata = [dict(id=e.get("id"), title=e.get("title"), scores_affected=e.get("scores_affected"), url=base + str(e.get("url", ""))) for e in ex["errata"]]
    all_match = all(s["fitness_matches"] for s in seats)

    if args.json:
        print(json.dumps(dict(match_id=args.match_id, base=base, rubric=rubric, match_rounds=report["matchTicks"], seats=seats, errata=errata), indent=1))
        return 0 if all_match else 1

    same = lambda ok: "same" if ok else "DIFFERS"
    print(f"Match {args.match_id}: rubric {rubric}, {len(seats)} seats, {report['matchTicks']} rounds, read from {base}\n")
    width = max([len(s["name"]) for s in seats] + [4])
    print(f"{'':{width}}  {'final wealth':>27}  {'Fitness Index':>27}")
    print(f"{'seat':{width}}  {'from the log':>12} {'archive':>9}{'':6}  {'recomputed':>12} {'report':>9}")
    for s in seats:
        print(f"{s['name']:{width}}  {s['wealth']:>12.2f} {s['archived_wealth']:>9.2f} {same(s['wealth_matches']):>7}"
              f"  {s['fitness_index']:>12.1f} {s['published_fitness_index']:>9.1f} {same(s['fitness_matches']):>7}")
    print()
    print(f"Fitness Index: {sum(s['fitness_matches'] for s in seats)} of {len(seats)} seats match the published report.")
    print(f"Final wealth: {sum(s['wealth_matches'] for s in seats)} of {len(seats)} seats match the archive.")
    print("Taken on trust: each seat's message count, because direct messages are private.")
    if errata:
        print("\nErrata on this match. Read them before trusting a recomputation; one can name logged events that never took effect:")
        for e in errata:
            print(f"  {e['id']}: {e['title']} (published scores affected: {'yes' if e['scores_affected'] else 'no'}) {e['url']}")
    return 0 if all_match else 1


if __name__ == "__main__":
    sys.exit(main())