#!/usr/bin/env python3
"""Run rolling-origin forecasting jobs (instrument x model x distribution) and write one Parquet per job.

usage: run_benchmark.py --jobs results/manifests/jobs_<name>.json [--shard i/N] [--workers 8] [--out results/forecasts]
The job list is produced by scripts/make_jobs.py from the frozen configs. Each job writes
  <out>/<symbol>/<spec_id>.parquet            (forecasts, one row per origin)
  <out>/<symbol>/<spec_id>.fits.parquet       (refit attempts)
and a line in <out>/_runs.jsonl (git SHA, host, package versions, wall time). Existing outputs are skipped (resumable).
Nothing printed contains return values."""
from __future__ import annotations

import argparse
import json
import os
import platform
import resource
import socket
import subprocess
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))


def git_sha() -> str:
    try:
        return subprocess.check_output(["git", "-C", str(ROOT), "rev-parse", "--short", "HEAD"], text=True).strip()
    except Exception:  # noqa: BLE001
        return "unknown"


def run_one(job: dict, out: str) -> dict:
    import numpy as np
    import pandas as pd
    from wp46.estimation import JobSpec, run_job
    from wp46.preprocessing import load_panel
    js = JobSpec(**{k: v for k, v in job.items() if k in JobSpec.__dataclass_fields__})
    d = Path(out) / js.symbol
    d.mkdir(parents=True, exist_ok=True)
    fp = d / f"{js.spec_id}.parquet"
    if fp.exists():
        return {"spec_id": js.spec_id, "symbol": js.symbol, "skipped": True}
    t0 = time.perf_counter()
    ru0 = resource.getrusage(resource.RUSAGE_SELF)
    panel = load_panel(js.symbol)
    F, G = run_job(panel, js)
    F.to_parquet(fp, index=False)
    G.to_parquet(d / f"{js.spec_id}.fits.parquet", index=False)
    ru = resource.getrusage(resource.RUSAGE_SELF)
    return {"spec_id": js.spec_id, "symbol": js.symbol, "n_origins": int(len(F)), "n_fits": int(len(G)),
            "status_counts": F["status"].value_counts().to_dict() if len(F) else {},
            "fit_status_counts": G["status"].value_counts().to_dict() if len(G) else {},
            "wall_s": time.perf_counter() - t0, "cpu_s": (ru.ru_utime - ru0.ru_utime) + (ru.ru_stime - ru0.ru_stime), "maxrss_mb": ru.ru_maxrss / (1024 ** 2 if platform.system() == "Darwin" else 1024),
            "fit_time_total_s": float(G["time_s"].sum()) if len(G) else 0.0, "fc_time_total_s": float(F["time_fc_s"].sum()) if "time_fc_s" in F else 0.0}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--jobs", required=True)
    ap.add_argument("--shard", default="0/1")
    ap.add_argument("--workers", type=int, default=max(1, os.cpu_count() - 2))
    ap.add_argument("--out", default=str(ROOT / "results" / "forecasts"))
    a = ap.parse_args()
    jobs = json.loads(Path(a.jobs).read_text())
    i, n = map(int, a.shard.split("/"))
    jobs = [j for k, j in enumerate(jobs) if k % n == i]
    import arch, numpy, scipy, pandas
    meta = {"git_sha": git_sha(), "host": socket.gethostname(), "python": platform.python_version(), "arch": arch.__version__,
            "numpy": numpy.__version__, "scipy": scipy.__version__, "pandas": pandas.__version__, "machine": platform.machine(),
            "cpu_count": os.cpu_count(), "shard": a.shard, "n_jobs": len(jobs), "started": time.strftime("%Y-%m-%dT%H:%M:%S%z")}
    Path(a.out).mkdir(parents=True, exist_ok=True)
    log = Path(a.out) / f"_runs_{socket.gethostname()}.jsonl"   # per host: rsync-merge safe
    print(json.dumps(meta), flush=True)
    done = 0
    t0 = time.time()
    with ProcessPoolExecutor(a.workers) as ex:
        futs = {ex.submit(run_one, j, a.out): j for j in jobs}
        for f in as_completed(futs):
            j = futs[f]
            try:
                rec = f.result()
            except Exception as e:  # noqa: BLE001
                rec = {"spec_id": j.get("model_id"), "symbol": j.get("symbol"), "error": f"{type(e).__name__}: {str(e)[:200]}"}
            rec.update({"git_sha": meta["git_sha"], "host": meta["host"], "ts": time.strftime("%Y-%m-%dT%H:%M:%S%z")})
            with open(log, "a") as fh:
                fh.write(json.dumps(rec) + "\n")
            done += 1
            if done % 10 == 0 or "error" in rec:
                print(f"[{done}/{len(jobs)} {time.time()-t0:7.0f}s] {rec.get('symbol')} {rec.get('spec_id')} {rec.get('error','') or rec.get('wall_s','')}", flush=True)
    print("done", done, "jobs in", round(time.time() - t0), "s")


if __name__ == "__main__":
    main()