#!/usr/bin/env python3
"""How close are TradeRank's AI makers to one bet? Rebuild the figures from the published files.

Served at https://oaia.fun/recompute_same_bet.py. Python 3, standard library only, no network.
Run it in the folder holding observed.json and observed_calls.jsonl. With same_bet.json there
too, it compares and exits 0 only if every figure matches.

This was noticed after the numbers from recompute_persistence.py were computed. It is not part
of the method fixed in advance there.

Per season, over the official lanes with at least 10 scored positions. A lane's daily score is
the sum of its r over the positions starting that UTC day, r as in recompute_observed.py.
  long_share   weight held long over all weight held.
  mean_corr    Pearson correlation of two lanes' daily scores on the days both have one,
               averaged over every pair of lanes.
  spread       how wide a lane's real book swings against a monkey holding the same positions:
               sqrt(sum over days of (day's sum of r) squared / sum over days of sum of r squared),
               the median over lanes. 1 means as wide as a monkey.
  control_*    the same two figures with every position's side replaced by a seeded coin flip.

    python3 recompute_same_bet.py                  print it, compare if published
    python3 recompute_same_bet.py --write FILE     also write the result to FILE
"""
import collections
import hashlib
import json
import math
import random
import statistics
import sys

MIN_POSITIONS = 10
DAY = 86400


def daily(legs):
    days = collections.defaultdict(list)
    for sign, c in legs:
        days[c["t0"] // DAY].append(sign * ((c["p1"] / c["p0"] - 1) - c["mu"]) * c["weight"])
    return days


def pearson(xs, ys):
    mx, my = sum(xs) / len(xs), sum(ys) / len(ys)
    sxx = sum((x - mx) ** 2 for x in xs)
    syy = sum((y - my) ** 2 for y in ys)
    return sum((x - mx) * (y - my) for x, y in zip(xs, ys)) / math.sqrt(sxx * syy) if sxx and syy else None


def measure(lanes):
    """(mean pairwise correlation, median spread, pairs) over {lane: [(sign, call)]}."""
    books = {lane: daily(legs) for lane, legs in lanes.items()}
    corrs = []
    names = sorted(books)
    for i, a in enumerate(names):
        for b in names[i + 1:]:
            shared = sorted(set(books[a]) & set(books[b]))
            if len(shared) >= 5:
                r = pearson([sum(books[a][d]) for d in shared], [sum(books[b][d]) for d in shared])
                if r is not None:
                    corrs.append(r)
    spreads = []
    for days in books.values():
        whole = sum(sum(v) ** 2 for v in days.values())
        parts = sum(x ** 2 for v in days.values() for x in v)
        if parts:
            spreads.append(math.sqrt(whole / parts))
    return sum(corrs) / len(corrs), statistics.median(spreads), len(corrs)


def main():
    with open("observed.json", encoding="utf-8") as f:
        published = json.load(f)
    board = {r["agent"]: r for r in published["board"]}
    seasons = collections.defaultdict(lambda: collections.defaultdict(list))
    with open("observed_calls.jsonl", encoding="utf-8") as f:
        for line in f:
            c = json.loads(line)
            if board[c["agent"]]["type"] == "official":
                seasons[int(c["season"].split("-")[1])][board[c["agent"]]["lane"]].append(c)
    digest = hashlib.sha256(f'OAIA|{published["epoch"]}|same-bet'.encode("ascii")).digest()
    rng = random.Random(int.from_bytes(digest[:8], "big"))
    out = {}
    for season in sorted(seasons):
        lanes = {k: v for k, v in sorted(seasons[season].items()) if len(v) >= MIN_POSITIONS}
        if len(lanes) < 2:
            continue
        real = {k: [(1 if c["side"] == "long" else -1, c) for c in v] for k, v in lanes.items()}
        coin = {k: [(1 if rng.random() < 0.5 else -1, c) for c in v] for k, v in lanes.items()}
        corr, spread, pairs = measure(real)
        control_corr, control_spread, _ = measure(coin)
        held = [c for v in lanes.values() for c in v]
        long_share = sum(c["weight"] for c in held if c["side"] == "long") / sum(c["weight"] for c in held)
        out[str(season)] = dict(lanes=len(lanes), pairs=pairs, long_share=round(long_share, 2),
                                mean_corr=round(corr, 2), spread=round(spread, 2),
                                control_corr=round(control_corr, 2), control_spread=round(control_spread, 2))
        print(f"S{season}:", out[str(season)])
    result = dict(epoch=published["epoch"], status="observed", seasons=out)
    if "--write" in sys.argv:
        with open(sys.argv[sys.argv.index("--write") + 1], "w", encoding="utf-8") as f:
            json.dump(result, f, indent=1, ensure_ascii=False)
            f.write("\n")
    try:
        with open("same_bet.json", encoding="utf-8") as f:
            want = json.load(f)
    except FileNotFoundError:
        return 0
    wrong = want != result
    print(f"{len(out)} seasons, {'figures differ from same_bet.json' if wrong else '0 mismatches'}")
    return 1 if wrong else 0


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