# /// script
# requires-python = "==3.12.*"
# dependencies = [
#   "prophet==1.1.6",
#   # Pinned, and the pins are load-bearing rather than cautious:
#   #  - cmdstanpy 1.3 breaks prophet 1.1.6's backend loader, and it fails as
#   #    "'Prophet' object has no attribute 'stan_backend'" — the real error is
#   #    swallowed and the traceback points at a logging line.
#   #  - pandas 3 breaks prophet's internal reindexing ("cannot reindex on an
#   #    axis with duplicate labels").
#   # Both surface per-series and would silently become a column of NaN, i.e. a
#   # benchmark where the competitor never ran.
#   "cmdstanpy==1.2.4",
#   "pandas==2.2.3",
#   "numpy<2.3",
# ]
# ///
"""
TimesFM zero-shot vs Prophet vs seasonal-naive on the M4 competition data.

Run it yourself, either against the hosted API or with no account at all:

    # Runs TimesFM on your own machine. No key, no signup. ~800 MB of weights
    # on first run; --device mps on Apple silicon, cuda on an NVIDIA box.
    uv run --with 'timesfm[torch]==2.0.2' benchmarks/run.py --timesfm local

    # Against the hosted API instead
    VECTIME_API_KEY=vct_live_… uv run benchmarks/run.py --freq hourly weekly

Why this exists: docs/marketing/positioning.md says accuracy claims must trace
to something reproducible, and that Google's numbers are Google's. This is the
number we are allowed to quote, and the point of publishing the script is that a
stranger can disagree with it by rerunning it.

Design notes, because they are the parts a reader should argue with:

- **Seasonal-naive is in here to be beaten.** It is the M4 baseline and it is
  embarrassingly hard to beat on some frequencies. A comparison that omits it is
  a comparison designed to flatter whoever commissioned it.
- **Results are reported per frequency, never pooled.** Pooling lets a win on
  one frequency bury a loss on another, and the losses are the useful part.
- **Prophet's interval is 80% by default and our band is q10–q90**, which is
  also 80%. So the coverage column compares like with like: what fraction of the
  held-out points actually landed inside the interval each method promised. A
  method that is accurate but overconfident is not safe to plan with, and that
  is invisible in an error metric.
- Sampling is seeded, so two runs of the same command compare the same series.
"""

import argparse
import io
import json
import logging
import os
import sys
import time
import urllib.request
import warnings
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import pandas as pd

warnings.filterwarnings("ignore")
# Prophet fits one Stan model per series and cmdstanpy logs two lines for each,
# which buries the progress output under thousands of lines on a real run.
os.environ.setdefault("CMDSTANPY_LOG_LEVEL", "CRITICAL")
logging.getLogger("cmdstanpy").disabled = True
logging.getLogger("prophet").setLevel(logging.CRITICAL)

ROOT = Path(__file__).resolve().parents[1]
CACHE = ROOT / "tmp" / "m4"
OUT = ROOT / "benchmarks" / "results"
M4 = "https://raw.githubusercontent.com/Mcompetitions/M4-methods/master/Dataset"

API = os.environ.get("VECTIME_API", "https://api.vectime.cloud")
KEY = os.environ.get("VECTIME_API_KEY")

# M4's official horizons and seasonal periods. Not ours to choose: using a
# different horizon than the competition did would make these numbers
# incomparable to every published M4 result, which is most of their value.
SPEC = {
    "hourly": {"horizon": 48, "period": 24, "file": "Hourly"},
    "weekly": {"horizon": 13, "period": 1, "file": "Weekly"},
    "daily": {"horizon": 14, "period": 1, "file": "Daily"},
    "monthly": {"horizon": 18, "period": 12, "file": "Monthly"},
}
# Prophet needs real timestamps to place its seasonalities, so the M4 start
# dates are used rather than invented ones.
FREQ_ALIAS = {"hourly": "h", "weekly": "W", "daily": "D", "monthly": "MS"}


def die(msg: str):
    sys.exit(f"\n  {msg}\n")


def fetch(url: str, dest: Path) -> Path:
    """Download once and keep it — the M4 files are large and immutable."""
    if dest.exists():
        return dest
    dest.parent.mkdir(parents=True, exist_ok=True)
    print(f"  downloading {dest.name} …", flush=True)
    with urllib.request.urlopen(url) as r, open(dest, "wb") as f:
        f.write(r.read())
    return dest


def load(freq: str, n: int, seed: int):
    """Return [(series_id, train_values, test_values, start_timestamp)]."""
    name = SPEC[freq]["file"]
    train = pd.read_csv(fetch(f"{M4}/Train/{name}-train.csv", CACHE / f"{name}-train.csv"))
    test = pd.read_csv(fetch(f"{M4}/Test/{name}-test.csv", CACHE / f"{name}-test.csv"))
    info = pd.read_csv(fetch(f"{M4}/M4-info.csv", CACHE / "M4-info.csv"))

    starts = dict(zip(info["M4id"], info["StartingDate"]))
    train = train.set_index("V1")
    test = test.set_index("V1")

    ids = sorted(set(train.index) & set(test.index))
    rng = np.random.default_rng(seed)
    picked = sorted(rng.choice(ids, size=min(n, len(ids)), replace=False).tolist())

    out = []
    for sid in picked:
        y = train.loc[sid].dropna().to_numpy(dtype=float)
        t = test.loc[sid].dropna().to_numpy(dtype=float)
        if len(y) < 2 * SPEC[freq]["period"] + 10 or len(t) == 0:
            continue
        # A handful of M4 start dates are unparseable; those series are skipped
        # rather than given a made-up one, which would quietly change what
        # Prophet is being asked to model.
        try:
            start = pd.to_datetime(starts.get(sid), dayfirst=True)
        except Exception:
            continue
        if pd.isna(start):
            continue
        out.append((sid, y, t, start))
    return out


# --------------------------------------------------------------------- metrics

def smape(actual, pred):
    denom = np.abs(actual) + np.abs(pred)
    # Where both are zero the term is defined as zero, not as a division error.
    return float(200.0 * np.mean(np.divide(np.abs(actual - pred), denom, out=np.zeros_like(denom), where=denom != 0)))


def mase(actual, pred, train, period):
    scale = np.mean(np.abs(train[period:] - train[:-period])) if len(train) > period else np.nan
    if not np.isfinite(scale) or scale == 0:
        return float("nan")
    return float(np.mean(np.abs(actual - pred)) / scale)


def coverage(actual, lo, hi):
    """Share of held-out points inside the interval the method promised."""
    if lo is None or hi is None:
        return float("nan")
    return float(np.mean((actual >= lo) & (actual <= hi)))


# --------------------------------------------------------------------- methods

def seasonal_naive(y, horizon, period):
    if period <= 1:
        return np.repeat(y[-1], horizon)
    reps = int(np.ceil(horizon / period))
    return np.tile(y[-period:], reps)[:horizon]


def run_prophet(y, horizon, start, freq):
    from prophet import Prophet

    dates = pd.date_range(start=start, periods=len(y), freq=FREQ_ALIAS[freq])
    df = pd.DataFrame({"ds": dates, "y": y})
    # interval_width=0.8 is Prophet's default and matches q10–q90 exactly, which
    # is the only reason the coverage column is a fair comparison.
    m = Prophet(interval_width=0.8)
    m.fit(df)
    future = m.make_future_dataframe(periods=horizon, freq=FREQ_ALIAS[freq], include_history=False)
    fc = m.predict(future)
    return (
        fc["yhat"].to_numpy(),
        fc["yhat_lower"].to_numpy(),
        fc["yhat_upper"].to_numpy(),
    )


def api_limits():
    req = urllib.request.Request(f"{API}/v1/models", headers={"Authorization": f"Bearer {KEY}"})
    with urllib.request.urlopen(req, timeout=30) as r:
        body = json.load(r)
    model = body["models"][0]
    return model, body.get("plan", "unknown")


def forecast_batch(series_list, horizon):
    """One /v1/forecast call for a batch, retrying a 429 on Retry-After."""
    payload = json.dumps({"series": [s.tolist() for s in series_list], "horizon": horizon, "quantiles": True}).encode()
    for attempt in range(6):
        req = urllib.request.Request(
            f"{API}/v1/forecast",
            data=payload,
            headers={"Authorization": f"Bearer {KEY}", "Content-Type": "application/json"},
        )
        try:
            with urllib.request.urlopen(req, timeout=300) as r:
                body = json.load(r)
            if body.get("stub"):
                die("the API returned a STUB forecast. Benchmarking the seasonal-naive fallback\n"
                    "  against seasonal-naive is not a result. Unset TIMESFM_STUB.")
            return body["forecasts"]
        except urllib.error.HTTPError as e:
            detail = e.read().decode()[:300]
            if e.code == 429 and attempt < 5:
                wait = int(e.headers.get("Retry-After") or 2) + 1
                print(f"    rate limited, waiting {wait}s", flush=True)
                time.sleep(wait)
                continue
            die(f"/v1/forecast returned {e.code}: {detail}")
    die("gave up after repeated rate limiting")


def api_backend():
    """Hosted path: /v1/models states what this key is actually allowed to do."""
    if not KEY:
        die("VECTIME_API_KEY is not set. Create a key at https://vectime.cloud/app/keys\n"
            "  (the free tier is enough), or pass --timesfm local to run the model\n"
            "  yourself with no account and no key at all.")
    model, plan = api_limits()
    info = {
        "label": f"api/{plan}",
        "id": model["id"],
        "where": f"api {API} · plan {plan}",
        "max_series": model["max_series_per_request"],
        "max_context": model["max_context"],
        "max_horizon": model["max_horizon"],
    }
    return info, forecast_batch


def local_backend(device: str, context: int):
    """In-process path, so this benchmark is reproducible without an account.

    The service's own wrapper is imported rather than reimplemented here. It
    carries two corrections that are not obvious and would silently produce
    wrong numbers if reinvented: TimesFM's quantile array puts the *mean* at
    index 0 (q50 is index 5, not 0), and on mps the batch has to be padded to
    global_batch_size in float32 because the library's internal padding uses
    float64 literals, which mps cannot represent at all.
    """
    svc = ROOT / "services" / "inference"
    if not (svc / "app" / "forecaster.py").exists():
        die("--timesfm local needs services/inference/ in the tree. Use the\n"
            "  hosted path instead, or clone the full repository.")

    # app.settings constructs its Settings at import time, so the environment
    # has to be set *before* the import. Doing it after silently gives you the
    # service defaults (4096 context) instead of the flags passed here, which
    # would quietly make these numbers incomparable to the published table.
    os.environ.update({
        "TIMESFM_DEVICE": device,
        "TIMESFM_MAX_CONTEXT": str(context),
        "TIMESFM_MAX_HORIZON": "512",
        "TIMESFM_STUB": "0",
    })
    sys.path.insert(0, str(svc))
    try:
        from app.forecaster import Forecaster
        from app.settings import Settings
    except ImportError as e:
        die(f"could not import the inference wrapper: {e}\n"
            "  Install the model extra:\n"
            "    uv run --with 'timesfm[torch]==2.0.2' benchmarks/run.py --timesfm local")

    cfg = Settings()
    print(f"\n  loading {cfg.model_repo} on device={device} …")
    print("  first run downloads ~800 MB of weights", flush=True)
    fc = Forecaster(cfg)
    # torch and timesfm are imported inside load(), not at the top of the
    # wrapper, so a missing model extra surfaces here rather than at the import
    # above. The wrapper also log.exception()s before re-raising, which prints a
    # traceback nobody needs for "you did not install it" — hence the silencing.
    logging.getLogger("inference.forecaster").disabled = True
    try:
        fc.load()
    except ImportError as e:
        die(f"{e}\n"
            "  The model extra is not installed. Either:\n"
            "    uv run --with 'timesfm[torch]==2.0.2' benchmarks/run.py --timesfm local\n"
            "  or use the hosted path with a free-tier key instead.")
    except Exception as e:
        die(f"loading {cfg.model_repo} on device={device} failed:\n"
            f"  {type(e).__name__}: {e}\n"
            "  --device cpu is the slow-but-always-available fallback.")

    def predict(series_list, horizon):
        out = fc.forecast([x.tolist() for x in series_list], horizon, want_quantiles=True)
        if out.get("stub"):
            die("the local model returned a STUB forecast. Benchmarking the\n"
                "  seasonal-naive fallback against seasonal-naive is not a result.")
        # The wrapper returns batch-major arrays; the API returns one dict per
        # series. Reshape here so the scoring loop is identical for both.
        q = out.get("quantiles") or {}
        return [
            {"point": out["point"][i], "quantiles": {k: v[i] for k, v in q.items()}}
            for i in range(len(series_list))
        ]

    info = {
        "label": f"local/{device}",
        "id": cfg.model_name,
        "where": f"local · {cfg.model_repo} · device {device}",
        "max_series": cfg.per_core_batch_size,
        "max_context": cfg.max_context,
        "max_horizon": cfg.max_horizon,
    }
    return info, predict


# ------------------------------------------------------------------------ main

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--freq", nargs="+", default=["hourly", "weekly"], choices=sorted(SPEC))
    ap.add_argument("--series", type=int, default=100, help="series sampled per frequency")
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--skip-prophet", action="store_true")
    ap.add_argument("--timesfm", choices=("api", "local"), default="api",
                    help="api = hosted /v1/forecast (needs a key); "
                         "local = load the weights in this process (needs none)")
    ap.add_argument("--device", choices=("auto", "cpu", "mps", "cuda"), default="auto",
                    help="--timesfm local only")
    ap.add_argument("--context", type=int, default=1024,
                    help="--timesfm local only: context window, defaulting to the "
                         "1,024 the published table was run at. The hosted path "
                         "takes this from the plan instead.")
    ap.add_argument("--summarise", nargs="+", metavar="CSV",
                    help="re-summarise existing per-series CSVs instead of running anything")
    args = ap.parse_args()

    # Re-reading saved runs needs no API key and no model: the expensive part is
    # already on disk, and combining separate frequency runs into one table is
    # the normal way this gets used.
    if args.summarise:
        frames = [pd.read_csv(f) for f in args.summarise]
        summarise(pd.concat(frames, ignore_index=True), OUT / "summary.csv")
        return

    if args.timesfm == "local":
        info, predict = local_backend(args.device, args.context)
    else:
        info, predict = api_backend()

    max_series, max_context = info["max_series"], info["max_context"]
    print(f"\n  {info['where']} · model {info['id']}")
    print(f"  batch {max_series} · context {max_context} · horizon cap {info['max_horizon']}\n")

    OUT.mkdir(parents=True, exist_ok=True)
    rows = []

    for freq in args.freq:
        horizon, period = SPEC[freq]["horizon"], SPEC[freq]["period"]
        if horizon > info["max_horizon"]:
            print(f"  {freq}: horizon {horizon} exceeds the cap {info['max_horizon']} — skipped")
            continue

        data = load(freq, args.series, args.seed)
        print(f"  {freq}: {len(data)} series · horizon {horizon} · seasonal period {period}")

        # TimesFM, batched. Context is truncated to the plan's window; the
        # truncation is recorded so a reader knows the model did not see more
        # history than this.
        contexts = [y[-max_context:] for _, y, _, _ in data]
        preds, q10s, q90s = [], [], []
        for i in range(0, len(contexts), max_series):
            chunk = contexts[i : i + max_series]
            for fc in predict(chunk, horizon):
                preds.append(np.array(fc["point"]))
                q = fc.get("quantiles") or {}
                q10s.append(np.array(q["q10"]) if "q10" in q else None)
                q90s.append(np.array(q["q90"]) if "q90" in q else None)
            print(f"    timesfm {min(i + max_series, len(contexts))}/{len(contexts)}", flush=True)

        # Both paths replace NaN/Inf with 0.0 before returning, because the API
        # contract is JSON and JSON has no NaN. That is right for the API and
        # wrong here: a failed forecast would be scored as if it had predicted
        # zero. An all-zero forecast is possible but rare, so it is reported
        # rather than dropped — check the series before trusting the row.
        blank = sum(1 for p in preds if not np.any(p))
        if blank:
            print(f"    warning: {blank}/{len(preds)} forecasts are all-zero — "
                  "likely NaN/Inf cleaned to 0.0, and scored as a real forecast")

        for idx, (sid, y, t, start) in enumerate(data):
            actual = t[:horizon]
            h = len(actual)

            def record(method, pred, lo=None, hi=None):
                rows.append({
                    "freq": freq, "series": sid, "method": method,
                    "timesfm_backend": info["label"],
                    "smape": smape(actual, pred[:h]),
                    "mase": mase(actual, pred[:h], y, period),
                    "coverage": coverage(actual, lo[:h] if lo is not None else None,
                                         hi[:h] if hi is not None else None),
                })

            record("timesfm", preds[idx], q10s[idx], q90s[idx])
            record("seasonal_naive", seasonal_naive(y, horizon, period))

            if not args.skip_prophet:
                try:
                    p, lo, hi = run_prophet(y, horizon, start, freq)
                    record("prophet", p, lo, hi)
                except Exception as e:
                    # A Prophet failure is data, not a reason to drop the series
                    # from the other methods' numbers.
                    rows.append({"freq": freq, "series": sid, "method": "prophet",
                                 "timesfm_backend": info["label"],
                                 "smape": float("nan"), "mase": float("nan"),
                                 "coverage": float("nan"), "error": str(e)[:120]})
            if (idx + 1) % 25 == 0:
                print(f"    baselines {idx + 1}/{len(data)}", flush=True)

    df = pd.DataFrame(rows)
    stamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
    per_series = OUT / f"per-series-{stamp}.csv"
    df.to_csv(per_series, index=False)
    summarise(df, OUT / f"summary-{stamp}.csv")
    print(f"\n  per-series → {per_series.relative_to(ROOT)}")


def summarise(df: pd.DataFrame, dest: Path):
    """Per-frequency table plus the two things a mean alone would hide."""
    scored = df[df["mase"].notna()]

    # Median as well as mean, because one catastrophic fit moves a mean a long
    # way — on M4 Weekly a single Prophet series scored MASE 158 against a
    # median of 2.5, and reporting only the mean would misrepresent Prophet as
    # far worse than it typically is.
    summary = (
        scored.groupby(["freq", "method"])[["smape", "mase", "coverage"]]
        .agg(["mean", "median"])
        .round(3)
    )
    summary.columns = [f"{c}_{s}" for c, s in summary.columns]
    summary = summary.drop(columns=["coverage_median"]).reset_index().sort_values(["freq", "mase_median"])
    summary.to_csv(dest, index=False)
    print("\n" + summary.to_string(index=False))
    print(f"\n  summary    → {dest.relative_to(ROOT)}")

    # The headline number is a mean over series, and a mean can win while losing
    # on a third of the individual series. That rate is the honest answer to
    # "will this be better on my data", so it is printed rather than left in a
    # CSV for someone to find.
    print("\n  Per-series, how often TimesFM is beaten:")
    for freq in scored["freq"].unique():
        sub = scored[scored["freq"] == freq]
        tf = sub[sub["method"] == "timesfm"].set_index("series")["mase"]
        for method in sorted(set(sub["method"]) - {"timesfm"}):
            other = sub[sub["method"] == method].set_index("series")["mase"]
            common = tf.index.intersection(other.index)
            if not len(common):
                continue
            lost = int((tf[common] > other[common]).sum())
            print(f"    {freq:<9} loses to {method:<15} on {lost:>3}/{len(common)} series "
                  f"({lost / len(common) * 100:.0f}%)")


if __name__ == "__main__":
    main()
