"""Realistic bounds for Hub Dollars.

The V4 low/high scenarios put every assumption at its extreme at the same time
(cheapest price x fewest tokens x no post-training x zero for uncertain bases).
Here each family's quantity uncertainty is drawn independently, and the GPU price
level is one shared draw (the market price moves for everyone together).
The realistic range is the 10th-90th percentile of the simulated total.

Run from the Space folder: python3 scripts/realistic_bounds.py
"""
import json, random, re, statistics

SEED, N = 20261018, 20000
# $/GPU-hour or $M per 1e24 FLOP used by V4 (low, central, high)
V4 = {"H100": (1.5, 2.5, 4), "A100": (1, 2, 3), "H800": (1, 2, 3), "B200": (3, 5, 8), "FLOP": (0.8, 1.8, 4)}
# Realistic price level for large training runs, relative to central ($2.5/H100-hour):
# $2.0-3.0 per H100-hour, shared by all families.
PRICE = (0.8, 1.0, 1.2)
# FLOP-based families also depend on hardware utilization (central 40%, realistic 35-50%):
# 1e24 FLOP at 40% of H100 peak = 0.70M H100-hours, i.e. $1.76M at $2.5/h.
UTIL = (0.35, 0.40, 0.50)


def price_kind(f):
    m = f["method"]
    if "FLOP" in m[:40] or "1e24" in m[:20] or "6ND" in m[:12]:
        return "FLOP"
    for gpu in ("H100", "A100", "H800", "B200"):
        if re.search(rf"{gpu}[ -]?(h\b|GPU|heures)|heures {gpu}|{gpu} over", m):
            return gpu
    return "A100" if "GPU-hours" in m or "heures GPU" in m else "FLOP"


def quantity_range(f):
    """Quantity factors relative to central, with the V4 price effect removed."""
    lo, mid, hi = V4[price_kind(f)]
    c = f["central_musd"]
    q_lo = f["low_musd"] / (c * lo / mid)
    q_hi = f["high_musd"] / (c * hi / mid)
    return max(0.0, min(q_lo, 1.0)), max(q_hi, 1.0)


def main():
    fams = json.load(open("v4/families-v4.json"))
    rng = random.Random(SEED)
    q = [quantity_range(f) for f in fams]
    kinds = [price_kind(f) for f in fams]
    grades = [f.get("evidence_grade") or ("B" if "tokens reported" in f["method"] else "C") for f in fams]
    totals, per_fam, per_org = [], [[] for _ in fams], {}
    orgs = sorted({f["organization"] for f in fams})
    for _ in range(N):
        p = rng.triangular(PRICE[0], PRICE[2], PRICE[1])
        p_flop = p * UTIL[1] / rng.triangular(UTIL[0], UTIL[2], UTIL[1])
        u_shared = rng.random()
        tot, org_tot = 0.0, dict.fromkeys(orgs, 0.0)
        for i, f in enumerate(fams):
            # two-piece draw: the V4 central value stays each family's median.
            # Partly assumed families share one draw: their token guesses come from the same prior.
            u = u_shared if grades[i] == "C" else rng.random()
            qf = q[i][0] + (1 - q[i][0]) * (2 * u) ** 0.5 if u < 0.5 else q[i][1] - (q[i][1] - 1) * (2 * (1 - u)) ** 0.5
            v = f["central_musd"] * (p_flop if kinds[i] == "FLOP" else p) * qf
            per_fam[i].append(v); org_tot[f["organization"]] += v; tot += v
        totals.append(tot)
        for o in orgs:
            per_org.setdefault(o, []).append(org_tot[o])

    pct = lambda xs, k: statistics.quantiles(xs, n=100)[k - 1]
    out_f = [{"organization": f["organization"], "family": f["family"], "central_musd": f["central_musd"],
              "realistic_low_musd": round(pct(per_fam[i], 10), 3), "realistic_high_musd": round(pct(per_fam[i], 90), 3),
              "price_basis": price_kind(f), "quantity_factor_range": [round(q[i][0], 3), round(q[i][1], 3)]}
             for i, f in enumerate(fams)]
    out_o = [{"organization": o, "central_musd": round(sum(f["central_musd"] for f in fams if f["organization"] == o), 3),
              "realistic_low_musd": round(pct(per_org[o], 10), 3), "realistic_high_musd": round(pct(per_org[o], 90), 3)}
             for o in orgs]
    summary = {
        "method": "Monte Carlo, seed %d, %d draws. Shared GPU price level triangular(%.1f, %.1f, %.1f) x central prices "
                  "($2-3 per H100-hour); FLOP-based families also divide by a shared utilization draw triangular(%.2f, %.2f, %.2f). "
                  "Per-family quantity factor: two-piece triangular between the V4 low and high quantities (V4 price effect removed), "
                  "median 1; partly assumed (grade C) families share one draw, others are independent. Realistic range = 10th-90th percentile of the total." % (SEED, N, *PRICE, *UTIL),
        "central_musd": round(sum(f["central_musd"] for f in fams), 3),
        "median_musd": round(statistics.median(totals), 3),
        "realistic_low_musd": round(pct(totals, 10), 3),
        "realistic_high_musd": round(pct(totals, 90), 3),
        "extreme_low_musd": round(sum(f["low_musd"] for f in fams), 3),
        "extreme_high_musd": round(sum(f["high_musd"] for f in fams), 3),
        "documented_input_central_musd": round(sum(f.get("documented_input_central_musd") or 0 for f in fams), 3),
    }
    # Monthly cumulative bounds: each family's draws booked at its first or latest release month
    months = [r["month"] for r in json.load(open("timeline/cumulative.json"))]
    win = {w["organization"] + "|" + w["family"]: w for w in json.load(open("timeline/release-windows.json"))}
    monthly = []
    by_mode = {}
    for mode, key in (("early", "first_month"), ("late", "last_month")):
        book = [win[f["organization"] + "|" + f["family"]][key] for f in fams]
        run = [0.0] * N
        rows = []
        for m in months:
            for i in [i for i, b in enumerate(book) if b == m]:
                run = [r + v for r, v in zip(run, per_fam[i])]
            rows.append((pct(run, 10), pct(run, 90)) if any(run) else (0.0, 0.0))
        by_mode[mode] = rows
    for j, m in enumerate(months):
        monthly.append({"month": m,
                        "realistic_low_early_musd": round(by_mode["early"][j][0], 3), "realistic_high_early_musd": round(by_mode["early"][j][1], 3),
                        "realistic_low_late_musd": round(by_mode["late"][j][0], 3), "realistic_high_late_musd": round(by_mode["late"][j][1], 3)})
    json.dump(monthly, open("timeline/cumulative-realistic.json", "w"), indent=1)
    with open("timeline/cumulative-realistic.csv", "w") as fh:
        fh.write(",".join(monthly[0]) + "\n")
        fh.writelines(",".join(str(v) for v in r.values()) + "\n" for r in monthly)

    json.dump(out_f, open("v4/families-realistic.json", "w"), indent=1)
    json.dump(out_o, open("v4/organizations-realistic.json", "w"), indent=1)
    json.dump(summary, open("v4/summary-realistic.json", "w"), indent=1)
    print(json.dumps(summary, indent=1))


if __name__ == "__main__":
    main()
