#!/usr/bin/env python3
"""Does last season's best AI trader stay best? Rebuild the answer from the published files.

Served at https://oaia.fun/recompute_persistence.py. Python 3, standard library only, no
network. Run it in the folder holding observed.json and observed_calls.jsonl (both served
beside it, TradeRank's AI traders, crypto positions, scored at https://oaia.fun/ai-traders).
With persistence.json there too, it compares and exits 0 only if every figure matches.

THE METHOD, written and fixed before any season-level number was computed:

The unit is a lane, one maker's seat at TradeRank (openai, anthropic, ...). The model in a
lane changes between seasons. Reverse lanes are left out: each mirrors another lane, so
keeping them would count one bet twice.

1. Score. For each season and lane, take the lane's positions in (snapshot_ms, coin, side)
   order, r = side x ((p1 / p0 - 1) - mu) x weight, and 2,048 monkeys that flip a coin for
   each side, seeded sha256("OAIA|TRADERANK-2026.07.1-S<season>|<lane>"). The percentile is
   where the lane's sum of r sits among theirs, ties counting half. Under 10 positions in a
   season, the lane is not ranked in that season (the board's own proving rule).
2. Pair. For each two consecutive seasons, Spearman's rho of the percentiles across the
   lanes ranked in both, ties sharing their average rank.
3. The one headline number. The mean of rho over pairs with at least 5 lanes, weighted by
   lanes. A thinner pair is shown and not pooled.
4. Chance. 20,000 shuffles. Each one, inside each season on its own, deals that season's
   percentiles back to its lanes at random, then recomputes the headline number. The two-sided
   p is (1 + shuffles with |rho| at or above ours) / 20,001. The bar is the shuffles' 95th
   percentile of |rho|.
5. Shown beside it, with no claim beyond the counts:
   a. each season's best lane, and where it lands the next season;
   b. how many lanes above the median of their pair stay above it, against the shuffles' mean;
   c. steps 1 to 4 again on single models, seeded by agent name, for a model ranked in both
      seasons, so no version change sits inside the comparison;
   d. how many lane-seasons sit below 50.
6. Power. 2,000 simulated worlds at each of three true correlations (0.3, 0.5, 0.7) from one
   season to the next, on the same lanes: the share of worlds that clear the bar.

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

COHORT = 2048
MIN_POSITIONS = 10
MIN_PAIR = 5
SHUFFLES = 20000
WORLDS = 2000
TRUE_RHOS = (0.3, 0.5, 0.7)


def rng_for(text):
    digest = hashlib.sha256(text.encode("ascii")).digest()
    return random.Random(int.from_bytes(digest[:8], "big"))


def score(legs, seed):
    """(positions, monkeys below, ties, percentile) for legs against their own cohort."""
    legs = sorted(legs, key=lambda c: (c["snapshot_ms"], c["coin"], c["side"]))
    rets = [(1 if c["side"] == "long" else -1) * ((c["p1"] / c["p0"] - 1) - c["mu"]) * c["weight"]
            for c in legs]
    rng = rng_for(seed)
    own = sum(rets)
    below = ties = 0
    for _ in range(COHORT):
        monkey = sum((1.0 if rng.random() < 0.5 else -1.0) * r for r in rets)
        below += monkey < own
        ties += monkey == own
    return len(rets), below, ties, 100 * (below + 0.5 * ties) / COHORT


def ranks(values):
    """Ranks from 1 for the lowest value; equal values share their average rank."""
    order = sorted(range(len(values)), key=values.__getitem__)
    out = [0.0] * len(values)
    i = 0
    while i < len(order):
        j = i
        while j + 1 < len(order) and values[order[j + 1]] == values[order[i]]:
            j += 1
        for k in range(i, j + 1):
            out[order[k]] = (i + j) / 2 + 1
        i = j + 1
    return out


def spearman(xs, ys):
    rx, ry = ranks(xs), ranks(ys)
    mean = (len(xs) + 1) / 2
    sxy = sum((a - mean) * (b - mean) for a, b in zip(rx, ry))
    sxx = sum((a - mean) ** 2 for a in rx)
    syy = sum((b - mean) ** 2 for b in ry)
    return sxy / math.sqrt(sxx * syy) if sxx and syy else 0.0


def above_median(values):
    middle = (len(values) + 1) / 2
    return [r > middle for r in ranks(values)]


def headline(pairs, values):
    """Lane-weighted mean rho and the repeat count, over pairs given as (k, k + 1, names)."""
    total = weight = repeat = 0
    for k, nxt, names in pairs:
        xs = [values[k][n] for n in names]
        ys = [values[nxt][n] for n in names]
        total += spearman(xs, ys) * len(names)
        weight += len(names)
        repeat += sum(a and b for a, b in zip(above_median(xs), above_median(ys)))
    return total / weight, repeat


def persistence(ranked, seed):
    """Steps 2 to 4 over {season: {name: percentile}}."""
    pairs = []
    for k in sorted(ranked):
        if k + 1 in ranked:
            pairs.append((k, k + 1, sorted(set(ranked[k]) & set(ranked[k + 1]))))
    rows = []
    for k, nxt, names in pairs:
        rho = spearman([ranked[k][n] for n in names], [ranked[nxt][n] for n in names]) if len(names) > 1 else None
        rows.append(dict(season_from=k, season_to=nxt, names=len(names),
                         rho=None if rho is None else round(rho, 3), pooled=len(names) >= MIN_PAIR))
    pooled = [p for p in pairs if len(p[2]) >= MIN_PAIR]
    rho, repeat = headline(pooled, ranked)
    rng = rng_for(seed)
    used = sorted({k for p in pooled for k in p[:2]})
    extreme = repeats = 0
    abs_rhos = []
    for _ in range(SHUFFLES):
        dealt = {}
        for k in used:
            names = sorted(ranked[k])
            deal = [ranked[k][n] for n in names]
            for i in range(len(deal) - 1, 0, -1):
                j = int(rng.random() * (i + 1))
                deal[i], deal[j] = deal[j], deal[i]
            dealt[k] = dict(zip(names, deal))
        shuffled_rho, shuffled_repeat = headline(pooled, dealt)
        abs_rhos.append(abs(shuffled_rho))
        extreme += abs(shuffled_rho) >= abs(rho) - 1e-12
        repeats += shuffled_repeat
    abs_rhos.sort()
    of = sum(sum(above_median([ranked[k][n] for n in names])) for k, _, names in pooled)
    return dict(pairs=rows, rho=round(rho, 3), p=round((1 + extreme) / (1 + SHUFFLES), 4),
                bar=round(abs_rhos[int(0.95 * SHUFFLES)], 3), pair_points=sum(len(p[2]) for p in pooled),
                top_half_repeat=repeat, top_half_of=of, top_half_chance=round(repeats / SHUFFLES, 2)), pooled


def power(pooled, ranked, bar, seed):
    """The share of simulated worlds that clear the bar at a given true correlation."""
    rng = rng_for(seed)
    out = {}
    used = sorted({k for p in pooled for k in p[:2]})
    names = sorted({n for k in used for n in ranked[k]})
    for true_rho in TRUE_RHOS:
        hits = 0
        for _ in range(WORLDS):
            world = {}
            latent = {n: rng.gauss(0, 1) for n in names}
            for k in used:
                world[k] = {n: latent[n] for n in ranked[k]}
                latent = {n: true_rho * latent[n] + math.sqrt(1 - true_rho ** 2) * rng.gauss(0, 1)
                          for n in names}
            hits += abs(headline(pooled, world)[0]) >= bar
        out[f"{true_rho:.1f}"] = round(hits / WORLDS, 2)
    return out


def main():
    with open("observed.json", encoding="utf-8") as f:
        published = json.load(f)
    epoch = published["epoch"]
    board = {r["agent"]: r for r in published["board"]}
    by_lane = collections.defaultdict(list)
    by_model = collections.defaultdict(list)
    with open("observed_calls.jsonl", encoding="utf-8") as f:
        for line in f:
            c = json.loads(line)
            row = board[c["agent"]]
            if row["type"] != "official":
                continue
            season = int(c["season"].split("-")[1])
            by_lane[season, row["lane"]].append(c)
            by_model[season, c["agent"]].append(c)
    makers = {r["lane"]: r["provider"] for r in board.values() if r["type"] == "official"}

    # Models a lane ran that season, in the order it first held a position with each.
    def models_of(legs):
        first = {}
        for c in sorted(legs, key=lambda c: c["snapshot_ms"]):
            first.setdefault(board[c["agent"]]["model"], c["snapshot_ms"])
        return sorted(first, key=first.get)

    grid = []
    ranked = collections.defaultdict(dict)
    for (season, lane), legs in sorted(by_lane.items()):
        n, below, ties, pct = score(legs, f"OAIA|{epoch}-S{season}|{lane}")
        grid.append(dict(season=season, lane=lane, positions=n, beat=below, ties=ties,
                         percentile=round(pct, 1), models=models_of(legs)))
        if n >= MIN_POSITIONS:
            ranked[season][lane] = pct
    for cell in grid:
        table = ranked[cell["season"]]
        if cell["lane"] in table:
            cell["rank"] = 1 + sum(v > table[cell["lane"]] for v in table.values())
            cell["ranked"] = len(table)

    lanes, pooled = persistence(ranked, f"OAIA|{epoch}|persistence-lanes")
    cells = {(c["season"], c["lane"]): c for c in grid}
    best = []
    for season in sorted(ranked):
        top = max(sorted(ranked[season]), key=ranked[season].get)
        row = dict(season=season, lane=top, models=cells[season, top]["models"],
                   percentile=cells[season, top]["percentile"], positions=cells[season, top]["positions"],
                   ranked=len(ranked[season]))
        following = cells.get((season + 1, top))
        if season + 1 in ranked:
            row["next_positions"] = following["positions"] if following else 0
            row["next_models"] = following["models"] if following else []
            if following and "rank" in following:
                row.update(next_rank=following["rank"], next_ranked=following["ranked"],
                           next_percentile=following["percentile"])
            best.append(row)

    single = collections.defaultdict(dict)
    single_rows = []
    for (season, agent), legs in sorted(by_model.items()):
        n, below, ties, pct = score(legs, f"OAIA|{epoch}-S{season}|{agent}")
        if n >= MIN_POSITIONS:
            single[season][agent] = pct
            single_rows.append(dict(season=season, agent=agent, model=board[agent]["model"],
                                    positions=n, beat=below, ties=ties, percentile=round(pct, 1)))
    models, _ = persistence(single, f"OAIA|{epoch}|persistence-models")

    below_50 = {}
    for season in sorted(ranked):
        below_50[str(season)] = [sum(v < 50 for v in ranked[season].values()), len(ranked[season])]

    result = dict(
        epoch=epoch, status="observed", cohort=COHORT, min_positions=MIN_POSITIONS,
        min_pair=MIN_PAIR, shuffles=SHUFFLES, worlds=WORLDS, makers=makers, grid=grid,
        lanes=lanes, best=best, single_models=dict(models, rows=single_rows), below_50=below_50,
        power=power(pooled, ranked, lanes["bar"], f"OAIA|{epoch}|persistence-power"))
    result = json.loads(json.dumps(result))

    for p in lanes["pairs"]:
        print(f'S{p["season_from"]} to S{p["season_to"]}: rho {p["rho"]} on {p["names"]} lanes'
              f'{"" if p["pooled"] else " (thin, not pooled)"}')
    print(f'pooled rho {lanes["rho"]}, two-sided p {lanes["p"]}, bar {lanes["bar"]}, '
          f'top half repeats {lanes["top_half_repeat"]} of {lanes["top_half_of"]} '
          f'(chance {lanes["top_half_chance"]})')
    print(f'same model only: rho {models["rho"]}, p {models["p"]}, bar {models["bar"]} '
          f'on {models["pair_points"]} model pairs')
    print("power at true rho:", result["power"])
    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("persistence.json", encoding="utf-8") as f:
            want = json.load(f)
    except FileNotFoundError:
        return 0
    wrong = [k for k in sorted(set(want) | set(result)) if want.get(k) != result.get(k)]
    for k in wrong:
        print(f"{k}: rebuilt figures differ from persistence.json")
    print(f"{len(result)} sections, {len(wrong)} mismatches")
    return 1 if wrong else 0


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