"""Recompute the GIFT-Eval aggregates used in AI Gen Apps paper AGA-57.

https://aigenapps.com/papers/research/timestral-vs-timesfm-3

Reads the per-dataset results that every submission to the GIFT-Eval leaderboard
publishes, and prints:

  1. each model's MASE and CRPS relative to Seasonal Naive, aggregated by
     geometric mean over the benchmark's configurations (the figures on the
     Timestral and TimesFM-3 model cards);
  2. the average per-configuration CRPS rank, the leaderboard's other ordering;
  3. a paired comparison of two models, with a bootstrap interval and the split
     by domain, term and univariate or multivariate.

Usage, standard library only:

  git clone https://github.com/SalesforceAIResearch/gift-eval.git
  python3 gift-eval-recompute.py gift-eval/results Timestral-209M TimesFM-3

The bootstrap resamples configurations, which are not independent (one dataset
appears at several terms), so treat the interval as optimistic.
"""

import csv
import json
import math
import os
import random
import sys

CRPS = "eval_metrics/mean_weighted_sum_quantile_loss"
MASE = "eval_metrics/MASE[0.5]"


def load(results, model):
    rows = {}
    with open(os.path.join(results, model, "all_results.csv")) as f:
        for r in csv.DictReader(f):
            rows[r["dataset"]] = {
                "mase": float(r[MASE]),
                "crps": float(r[CRPS]),
                "domain": r["domain"],
                "variates": int(r["num_variates"]),
            }
    return rows


def config(results, model):
    path = os.path.join(results, model, "config.json")
    return json.load(open(path)) if os.path.exists(path) else {}


def geo_mean(xs):
    return math.exp(sum(math.log(x) for x in xs) / len(xs))


def main(results, a, b):
    base = load(results, "Seasonal_Naive")
    configs = sorted(base)
    models = {}
    for name in sorted(os.listdir(results)):
        try:
            rows = load(results, name)
        except (OSError, KeyError, ValueError):
            continue
        if set(rows) == set(base):
            models[name] = (config(results, name), rows)

    print(f"{len(configs)} configurations, {len(models)} complete submissions\n")
    print("Zero-shot or pretrained, no declared test leakage, by geometric mean CRPS")
    clean = [
        m for m, (cfg, _) in models.items()
        if cfg.get("testdata_leakage") == "No" and cfg.get("model_type") in ("zero-shot", "pretrained")
    ]
    scored = []
    for m in clean:
        rows = models[m][1]
        crps = geo_mean([rows[c]["crps"] / base[c]["crps"] for c in configs])
        mase = geo_mean([rows[c]["mase"] / base[c]["mase"] for c in configs])
        scored.append((crps, mase, m))
    for i, (crps, mase, m) in enumerate(sorted(scored)[:15], 1):
        print(f"{i:3d}  CRPS {crps:.3f}  MASE {mase:.3f}  {m}")

    print("\nSame set, by average CRPS rank across all complete submissions")
    rank = {m: 0.0 for m in models}
    for c in configs:
        for r, m in enumerate(sorted(models, key=lambda m: models[m][1][c]["crps"]), 1):
            rank[m] += r / len(configs)
    for i, m in enumerate(sorted(clean, key=rank.get)[:15], 1):
        print(f"{i:3d}  {rank[m]:6.2f}  {m}")

    ra, rb = models[a][1], models[b][1]
    random.seed(0)
    for metric in ("mase", "crps"):
        logs = [math.log(ra[c][metric] / rb[c][metric]) for c in configs]
        boots = sorted(
            math.exp(sum(random.choice(logs) for _ in logs) / len(logs)) for _ in range(10000)
        )
        wins = sum(x < 0 for x in logs)
        print(
            f"\n{metric.upper()} {a} / {b}: {math.exp(sum(logs) / len(logs)):.4f}, "
            f"95% interval [{boots[250]:.4f}, {boots[9750]:.4f}], {a} better on {wins} of {len(configs)}"
        )
        groups = {}
        for c, x in zip(configs, logs):
            groups.setdefault(("domain", ra[c]["domain"]), []).append(x)
            groups.setdefault(("term", c.split("/")[-1]), []).append(x)
            kind = "multivariate" if ra[c]["variates"] > 1 else "univariate"
            groups.setdefault(("shape", kind), []).append(x)
        for (kind, name), xs in sorted(groups.items()):
            print(f"   {kind:6s} {name:14s} n={len(xs):3d}  ratio {math.exp(sum(xs) / len(xs)):.3f}  wins {sum(x < 0 for x in xs)}")


if __name__ == "__main__":
    if len(sys.argv) != 4:
        sys.exit("usage: gift-eval-recompute.py <gift-eval/results> <model A> <model B>")
    main(*sys.argv[1:])