"""Reproduce a time-equivalent progress comparison from saved LLM Stats data.

Run: python analyze.py
Inputs preserve every listed model, score and announcement date, with source links.
This fits historical frontier records; it does not equate task difficulty.
"""
from __future__ import annotations

import csv
import hashlib
import json
import math
from datetime import date, timedelta
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import minimize_scalar

ROOT = Path(__file__).resolve().parent
(ROOT / "figures").mkdir(exist_ok=True)
ORIGIN = date(2026, 1, 1)
D0, D1 = .790, .823
MONTH_DAYS = 365.25 / 12


def frontier(rows):
    """At most one record per date; same-day ties do not add fit weight."""
    best = -math.inf
    result = []
    for row in sorted(rows, key=lambda r: (r["announcement_date"], -r["score"], r["model_id"])):
        if row["score"] > best:
            result.append(row)
            best = row["score"]
    return result


def xy(rows):
    return (np.array([(date.fromisoformat(r["announcement_date"]) - ORIGIN).days for r in rows], dtype=float),
            np.array([r["score"] for r in rows], dtype=float))


def linear(x, y):
    b, a = np.polyfit(x, y, 1)
    residuals = y - a - b*x
    return {"a": float(a), "b": float(b), "r2": float(1 - np.sum(residuals**2)/np.sum((y-y.mean())**2)),
            "rmse_pp": float(100*np.sqrt(np.mean(residuals**2)))}


def shifted_log(x, y):
    # Profile OLS over the time shift. Every log argument must be positive.
    # Optimize ln(first_date - tau) over a broad grid of brackets, then select
    # the least SSE. This avoids choosing an arbitrary calendar origin or
    # relying on one unexamined local optimum. Bounds are reported in results.
    def at_log_offset(q):
        tau = float(x.min() - math.exp(q))
        f = linear(np.log(x-tau), y)
        residuals = y - f["a"] - f["b"]*np.log(x-tau)
        return float(np.sum(residuals**2)), f, tau

    lo, hi = math.log(1e-4), math.log(36500)
    grid = np.linspace(lo, hi, 161)
    candidates = []
    for left, right in zip(grid[:-1], grid[1:]):
        opt = minimize_scalar(lambda q: at_log_offset(q)[0], bounds=(left, right), method="bounded",
                              options={"xatol": 1e-12})
        candidates.append((opt.fun, opt.x))
    _, q = min(candidates)
    sse, f, tau = at_log_offset(q)
    return {**f, "tau": tau, "sse": sse, "offset_days": math.exp(q),
            "offset_search_days": [1e-4, 36500], "parameters": 3}


def inverse_gap(f):
    return math.exp((D1-f["a"])/f["b"]) - math.exp((D0-f["a"])/f["b"])


data, records, manifest = {}, {}, {}
for slug in ["swe-bench-pro", "deepswe-1.1"]:
    raw = (ROOT/f"{slug}.csv").read_bytes()
    metadata = json.loads((ROOT/"source-metadata.json").read_text())[slug]
    with (ROOT/f"{slug}.csv").open(newline="") as source:
        data[slug] = {**metadata, "models": [{**row, "score": float(row["score"])} for row in csv.DictReader(source)]}
    records[slug] = frontier(data[slug]["models"])
    manifest[slug] = {"source": f"https://llm-stats.com/benchmarks/{slug}",
                      "rows": len(data[slug]["models"]), "records": len(records[slug]),
                      "csv_sha256": hashlib.sha256(raw).hexdigest(),
                      "source_updated_at": data[slug]["updated_at"],
                      "verified_count": data[slug]["statistics"]["verified_count"]}
    with (ROOT/f"{slug}-records.csv").open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=["announcement_date", "model_id", "model_name", "score", "publication_note"], lineterminator="\n")
        writer.writeheader()
        writer.writerows({k: r[k] for k in writer.fieldnames} for r in records[slug])

xp, yp = xy(records["swe-bench-pro"])
xd, yd = xy(records["deepswe-1.1"])
pro, deep = linear(xp, yp), shifted_log(xd, yd)
days = inverse_gap(deep)
equivalent = 100*pro["b"]*days
anchored_days = (xd[-1]-deep["tau"])*math.expm1((D1-D0)/deep["b"])

sensitivity = []
def add_case(name, duration, pro_slope=pro["b"], explanation=""):
    sensitivity.append({"method": name, "days": float(duration),
                        "equivalent_pp": float(100*pro_slope*duration), "explanation": explanation})

add_case("Full-history shifted log + linear", days, explanation="Primary model; fit the time origin and invert both score levels.")
add_case("Anchor DeepSWE curve to latest record", anchored_days,
         explanation="Same curvature, vertically shifted to pass through 79.0% on October 12.")
jan_log = linear(np.log(xd), yd)
add_case("Fix logarithmic origin at January 1", inverse_gap(jan_log),
         explanation="An arbitrary fixed origin changes the fitted curvature; poorer fit (R² 0.898).")
trimmed_log = shifted_log(xd[1:], yd[1:])
add_case("Omit earliest DeepSWE record", inverse_gap(trimmed_log), explanation="Seven records instead of eight.")
recent_deep = linear(xd[2:], yd[2:])
add_case("Use April–October DeepSWE linear trend", (D1-D0)/recent_deep["b"],
         explanation="Recent six DeepSWE records are also close to a line (R² 0.990).")
recent_pro = linear(xp[3:], yp[3:])
add_case("Use April–September Pro linear trend", days, recent_pro["b"],
         explanation="Only three Pro frontier records; its slope is lower than the full-history slope.")

holdouts = []
for count in [4, 5, 6, 7]:
    train_x, train_y, test_x, test_y = xd[:count], yd[:count], xd[count:], yd[count:]
    f = shifted_log(train_x, train_y)
    prediction = f["a"] + f["b"]*np.log(test_x-f["tau"])
    lf = linear(train_x, train_y)
    lp = lf["a"] + lf["b"]*test_x
    holdouts.append({"training_records": count, "held_out_records": len(test_x),
                     "shifted_log_rmse_pp": float(100*np.sqrt(np.mean((prediction-test_y)**2))),
                     "linear_rmse_pp": float(100*np.sqrt(np.mean((lp-test_y)**2)))})

results = {"snapshot_date": "2026-11-01", "retrieved_at_utc": "2026-11-01T22:26:28Z",
           "time_basis": "Model announcement_date in LLM Stats; not verified score-publication dates",
           "t_definition": "Days since 2026-01-01", "selection": "Strict new maxima, one per announcement date",
           "weighting": "Equal weight per record; no repeated daily plateau observations",
           "manifest": manifest, "pro_linear": pro, "deep_shifted_log": deep,
           "score_levels": [D0, D1], "equivalent_days": days, "equivalent_pro_pp": equivalent,
           "pro_pp_per_month": pro["b"]*100*MONTH_DAYS,
           "anchored_pro_pp": 100*pro["b"]*anchored_days,
           "deep_fitted_crossings_days": [deep["tau"]+math.exp((v-deep["a"])/deep["b"]) for v in [D0, D1]],
           "sensitivity": sensitivity, "rolling_holdout": holdouts,
           "exact_arithmetic": {"absolute_lead_pp": 100*(D1-D0), "relative_score_gain_pct": 100*(D1/D0-1),
                                "unsolved_share_reduction_pct": 100*(D1-D0)/(1-D0)}}
(ROOT/"results.json").write_text(json.dumps(results, indent=2)+"\n")

# Numerical consistency checks for the derivation, not predictive guarantees.
assert len(records["swe-bench-pro"]) == 6 and len(records["deepswe-1.1"]) == 8
assert np.all(np.diff(xd)>0) and np.all(np.diff(yd)>0)
for score in [D0, D1]:
    t = deep["tau"] + math.exp((score-deep["a"])/deep["b"])
    assert math.isclose(deep["a"]+deep["b"]*math.log(t-deep["tau"]), score, abs_tol=1e-12)
assert math.isclose(pro["b"]*days, equivalent/100, abs_tol=1e-12)

# Publication figures. Raw points, fitted curves and extrapolation remain distinct.
BG, FG, GRID = "#101318", "#e7ebf0", "#343b45"
PURPLE, BLUE, GREEN, MUTED = "#bd9cff", "#72b7ff", "#67dbb4", "#a8b1bf"
plt.rcParams.update({"figure.facecolor": BG, "axes.facecolor": BG, "savefig.facecolor": BG,
    "text.color": FG, "axes.labelcolor": MUTED, "xtick.color": MUTED, "ytick.color": MUTED,
    "font.family": "DejaVu Sans", "font.size": 11, "axes.edgecolor": GRID,
    "svg.fonttype": "none", "axes.titleweight": "bold"})

def format_axis(ax):
    ax.grid(axis="y", color=GRID, linewidth=.7)
    ax.spines[["top","right"]].set_visible(False)
    ax.xaxis.set_major_locator(mdates.MonthLocator(interval=2))
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%b %Y"))
    ax.set_ylabel("Reported score (%)")
    ax.set_ylim(0, 100)
    ax.set_yticks([0,20,40,60,80,100])
    ax.tick_params(axis="x", rotation=0)

def dates(values):
    return [ORIGIN+timedelta(days=float(v)) for v in values]

fig, axes = plt.subplots(1,2,figsize=(14,5.6))
for ax, slug, color, f, x, y in [(axes[0],"swe-bench-pro",BLUE,pro,xp,yp),
                                 (axes[1],"deepswe-1.1",PURPLE,deep,xd,yd)]:
    ax.set_title("SWE-bench Pro" if slug=="swe-bench-pro" else "DeepSWE v1.1",loc="left",pad=16,fontsize=17)
    all_x,all_y=xy(data[slug]["models"])
    ax.scatter(dates(all_x),100*all_y,s=19,color=MUTED,alpha=.22,zorder=1)
    ax.step(dates(x),100*y,where="post",color=color,alpha=.38,lw=1.5,zorder=2)
    ax.scatter(dates(x),100*y,s=58,color=color,edgecolor=BG,zorder=5)
    grid=np.linspace(x.min(),x.max(),400)
    pred=f["a"]+f["b"]*(grid if slug=="swe-bench-pro" else np.log(grid-f["tau"]))
    ax.plot(dates(grid),100*pred,color=color,lw=2.5,zorder=3)
    ax.set_xlim(date(2026,1,1),date(2026,11,1))
    format_axis(ax)
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%b"))
    kind="Linear" if slug=="swe-bench-pro" else "Shifted logarithmic"
    ax.text(.02,.05,f"{len(x)} frontier records · {kind} fit\nR² = {f['r2']:.3f}",transform=ax.transAxes,color=FG,fontsize=11)
fig.subplots_adjust(left=.065,right=.98,top=.84,bottom=.17,wspace=.23)
fig.text(.065,.95,"The actual leaderboard curves",fontsize=21,weight="bold")
fig.text(.065,.035,"2026 announcement dates · Faint dots: all entries · Solid dots: new records · Lines: fits · Snapshot: November 1, 2026",fontsize=10,color=MUTED)
for ext in ["svg","png"]:fig.savefig(ROOT/f"figures/frontiers.{ext}",dpi=180)
plt.close(fig)

fig,ax=plt.subplots(figsize=(12,5.1))
cut=xd[-1];t0,t1=results["deep_fitted_crossings_days"]
xx=np.linspace(155,cut,300); yy=deep["a"]+deep["b"]*np.log(xx-deep["tau"])
ax.plot(dates(xx),100*yy,color=PURPLE,lw=3,label="Fitted DeepSWE trend")
xx=np.linspace(cut,t1+9,200);yy=deep["a"]+deep["b"]*np.log(xx-deep["tau"])
ax.plot(dates(xx),100*yy,color=PURPLE,lw=3,ls="--",label="Extrapolation")
ax.scatter(dates(xd[3:]),100*yd[3:],s=70,color=PURPLE,edgecolor=BG,zorder=4)
ax.scatter(dates([t0,t1]),[79,82.3],s=85,color=GREEN,zorder=5)
for t,score,label in [(t0,79,"79.0%"),(t1,82.3,"82.3%")]:
    ax.vlines(dates([t]),69.5,score,color=GREEN,alpha=.5,lw=1,ls=":")
    ax.annotate(label,(mdates.date2num(dates([t])[0]),score),xytext=(-12,13),textcoords="offset points",color=GREEN,weight="bold",fontsize=13)
ax.annotate("",(mdates.date2num(dates([t0])[0]),70.5),(mdates.date2num(dates([t1])[0]),70.5),arrowprops={"arrowstyle":"<->","color":GREEN,"lw":2})
ax.text(mdates.date2num(dates([(t0+t1)/2])[0]),71.2,f"{days:.1f} days of fitted progress",ha="center",color=GREEN,fontsize=13)
ax.set_xlim(date(2026,6,1),dates([t1+22])[0]);format_axis(ax);ax.set_ylim(69,85);ax.set_yticks([70,75,79,82.3,85])
ax.set_title("Map the score gap to time, then apply the Pro trend",loc="left",pad=20,fontsize=17)
ax.legend(loc="upper left",facecolor=BG,edgecolor=GRID,labelcolor=FG,fontsize=10)
fig.subplots_adjust(left=.08,right=.96,top=.82,bottom=.2)
fig.text(.08,.04,f"{days:.1f} days × {pro['b']*100:.4f} Pro points/day = {equivalent:.1f} Pro points of time-equivalent progress",fontsize=13,color=FG,weight="bold")
for ext in ["svg","png"]:fig.savefig(ROOT/f"figures/equivalence.{ext}",dpi=180)
plt.close(fig)
print(json.dumps({"equivalent_pp":equivalent,"anchored_pp":results["anchored_pro_pp"],"days":days,
                  "pro":pro,"deep":deep,"checks":"passed"},indent=2))