"""Python timing for the same serving-shaped work as the Rust spike."""

import sys
import time
from pathlib import Path

import joblib
import numpy as np
import pandas as pd

REPO = Path(__file__).resolve().parents[3]  # repo root
sys.path.insert(0, str(REPO))
from src.ml import features  # noqa: E402

bundle = joblib.load(REPO / "models" / "flood_P.1.joblib")
stations = ["P.1"] + [u for u, _ in features.UPSTREAM_LEADS["P.1"]]
long = pd.concat(
    pd.read_csv(REPO / "models" / "cache" / f"{s}.csv.gz", parse_dates=["timestamp"])
    for s in stations
)
df = long[(long.timestamp >= "2026-07-30") & (long.timestamp < "2026-08-13")]
rain_all = pd.read_csv(REPO / "models" / "cache" / "rain_openmeteo.csv.gz", parse_dates=["timestamp"])
rain_all = rain_all.set_index("timestamp")["chiang_dao"].astype(float)

reps = 20
t = time.perf_counter()
for _ in range(reps):
    grid = features.make_hourly_grid(df)
    rain = rain_all.reindex(grid.observed.index)
    row = features.build_features(grid, "P.1", rain=rain)[bundle["feature_names"]].iloc[[-1]]
    out = {
        k: (h.predict_proba(row)[0, 1] if hasattr(h, "predict_proba") else h.predict(row)[0])
        for k, h in bundle["heads"].items()
    }
print(f"python: grid + features + 9 heads for the last hour: {(time.perf_counter() - t) / reps * 1000:.1f} ms per station")
t = time.perf_counter()
joblib.load(REPO / "models" / "flood_P.1.joblib")
print(f"python: bundle load {(time.perf_counter() - t) * 1000:.0f} ms")