#!/usr/bin/env python3
"""Reproduce the analysis in Section 4 of
"How long can an AI work on its own? Reach, reliability and the limits of measurement"
https://www.vlabsai.com/research/how-long-can-ai-work-alone

The script downloads METR's published results file for the Time Horizon 1.1
suite and computes, for every model released since January 2023:

  * the ratio of the 50% time horizon to the 80% time horizon,
  * the median, quartiles and range of that ratio,
  * the trend of the ratio over time (least squares on its logarithm),
  * the doubling time of each horizon over the frontier models,
  * the extrapolation table for higher success rates (Table 3).

It uses only the Python standard library. The data belong to METR and are not
redistributed here; the script reads them from METR's site each time it runs,
so results will change if METR revises the file.

    python3 analysis.py

Source of the data: https://metr.org/time-horizons/
"""
import datetime
import math
import re
import statistics
import urllib.request

DATA_URL = "https://metr.org/assets/benchmark_results_1_1.yaml"
SINCE = datetime.date(2023, 1, 1)
CEILING_MINUTES = 16 * 60  # METR: measurements above 16 hours are unreliable


def load(url=DATA_URL):
    """Parse the fields we need from METR's YAML without a YAML dependency."""
    request = urllib.request.Request(url, headers={"User-Agent": "vlabs-research-analysis"})
    text = urllib.request.urlopen(request, timeout=60).read().decode("utf-8")
    models, current, metric = {}, None, None
    for line in text.split("\n"):
        match = re.match(r"^  ([a-z0-9_.\-]+):\s*$", line)
        if match and not line.startswith("   "):
            current, metric = match.group(1), None
            models[current] = {}
            continue
        if current is None:
            continue
        match = re.match(r"^      (p50_horizon_length|p80_horizon_length):", line)
        if match:
            metric = match.group(1)
            continue
        match = re.match(r"^        estimate:\s*([0-9.eE+-]+)", line)
        if match and metric:
            models[current][metric] = float(match.group(1))
            metric = None
            continue
        match = re.match(r"^    release_date:\s*(\S+)", line)
        if match:
            models[current]["date"] = datetime.date.fromisoformat(match.group(1))
        match = re.match(r"^      is_sota:\s*(\S+)", line)
        if match:
            models[current]["sota"] = match.group(1) == "true"
    rows = []
    for name, values in models.items():
        if {"p50_horizon_length", "p80_horizon_length", "date"} <= values.keys():
            rows.append({"name": name, "p50": values["p50_horizon_length"], "p80": values["p80_horizon_length"],
                         "date": values["date"], "sota": values.get("sota", False)})
    return sorted(rows, key=lambda row: row["date"])


def least_squares(xs, ys):
    """Slope and correlation coefficient of y on x."""
    n = len(xs)
    mean_x, mean_y = sum(xs) / n, sum(ys) / n
    sxx = sum((x - mean_x) ** 2 for x in xs)
    syy = sum((y - mean_y) ** 2 for y in ys)
    sxy = sum((x - mean_x) * (y - mean_y) for x, y in zip(xs, ys))
    return sxy / sxx, sxy / math.sqrt(sxx * syy)


def main():
    rows = [row for row in load() if row["date"] >= SINCE]
    for row in rows:
        row["ratio"] = row["p50"] / row["p80"]

    print(f"Models released since {SINCE}: {len(rows)}\n")
    print(f"{'release':10}  {'model':40} {'50% (min)':>10} {'80% (min)':>10} {'ratio':>6}")
    for row in rows:
        print(f"{row['date']}  {row['name'][:40]:40} {row['p50']:10.1f} {row['p80']:10.1f} {row['ratio']:6.1f}")

    ratios = [row["ratio"] for row in rows]
    q1, median, q3 = statistics.quantiles(ratios, n=4)
    print(f"\nRatio of 50% to 80% horizon: median {median:.1f}, quartiles {q1:.1f} to {q3:.1f}, range {min(ratios):.1f} to {max(ratios):.1f}")

    years = [(row["date"] - SINCE).days / 365.25 for row in rows]
    slope, r = least_squares(years, [math.log(ratio) for ratio in ratios])
    print(f"Trend in the ratio: {100 * (math.exp(slope) - 1):+.0f}% a year (r = {r:.2f})")
    kept = [(year, ratio) for year, ratio, row in zip(years, ratios, rows) if row["ratio"] < max(ratios)]
    slope_without, _ = least_squares([year for year, _ in kept], [math.log(ratio) for _, ratio in kept])
    print(f"Trend without the largest ratio: {100 * (math.exp(slope_without) - 1):+.0f}% a year")

    frontier = [row for row in rows if row["sota"] and row["p50"] <= CEILING_MINUTES]
    days = [(row["date"] - SINCE).days for row in frontier]
    print(f"\nDoubling times over {len(frontier)} frontier models with a 50% horizon of at most 16 hours:")
    for key, label in (("p50", "50% horizon"), ("p80", "80% horizon")):
        slope, _ = least_squares(days, [math.log2(row[key]) for row in frontier])
        print(f"  {label}: {1 / slope:.0f} days")
    doubling = 131  # between the two fitted values
    print(f"Lag implied by the median ratio: {math.log2(median):.1f} doublings, about {math.log2(median) * doubling / 30.44:.0f} months")

    print("\nTask length at each success rate, extrapolated from a 16-hour 50% horizon (minutes):")
    exponent = math.log(median) / math.log(4)  # log-logistic slope implied by the median ratio
    print(f"{'success':>8} {'log-logistic':>13} {'const. hazard':>14} {'Weibull k=0.6':>14}")
    for p in (0.5, 0.8, 0.9, 0.95, 0.99):
        odds = p / (1 - p)
        log_logistic = CEILING_MINUTES / odds ** exponent
        hazard_fraction = math.log(1 / p) / math.log(2)
        print(f"{p:8.0%} {log_logistic:13.1f} {CEILING_MINUTES * hazard_fraction:14.1f} {CEILING_MINUTES * hazard_fraction ** (1 / 0.6):14.1f}")


if __name__ == "__main__":
    main()