#!/usr/bin/env python3
"""Thread scaling of ``simurg genotype`` (plan §18 P9 gate: ≥ 20 samples/min/core on GSA-size arrays, linear
scaling to 32 threads).

For every worker count the same N samples are genotyped (``--no-cohort-qc``); the per-sample phase is timed from
the run log ("genotyping … sample(s)" → "done"), so manifest loading is reported separately and does not blur
the scaling. Writes a JSON summary.

    python bench/scaling.py --bpm M.bpm --egt C.egt --idats DIR --out OUT --threads 1,2,4,8 [--samples N]
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import subprocess
import sys
import time
from pathlib import Path


def phase_seconds(log: Path) -> tuple[float, float]:
    """(manifest+setup seconds, per-sample phase seconds) from a genotype run log."""
    ts: dict[str, dt.datetime] = {}
    for line in log.read_text().splitlines():
        rec = json.loads(line)
        msg = rec.get("msg") or ""
        t = dt.datetime.fromisoformat(rec["ts"].replace("Z", "+00:00"))
        ts.setdefault("first", t)
        if msg.startswith("genotyping "):
            ts["start"] = t
        elif msg.startswith("done:"):
            ts["done"] = t
    return (ts["start"] - ts["first"]).total_seconds(), (ts["done"] - ts["start"]).total_seconds()


def subset(idats: Path, n: int, dest: Path) -> Path:
    grn = sorted(p for p in idats.rglob("*_Grn.idat*"))
    if n <= 0 or n >= len(grn):
        return idats
    dest.mkdir(parents=True, exist_ok=True)
    for g in grn[:n]:
        d = dest / g.parent.name
        d.mkdir(exist_ok=True)
        for p in (g, g.with_name(g.name.replace("_Grn.", "_Red."))):
            link = d / p.name
            if not link.exists():
                link.symlink_to(p.resolve())
    return dest


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--bpm", required=True)
    ap.add_argument("--egt", required=True)
    ap.add_argument("--idats", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--threads", default="1,2,4")
    ap.add_argument("--samples", type=int, default=0, help="use the first N samples (0 = all)")
    ap.add_argument("--simurg", default=str(Path(sys.executable).with_name("simurg")))
    ap.add_argument("--pure-python", action="store_true", help="run with SIMURG_PURE_PYTHON=1")
    a = ap.parse_args()
    a.out.mkdir(parents=True, exist_ok=True)
    idats = subset(a.idats, a.samples, a.out / "idats")
    n = len(list(idats.rglob("*_Grn.idat*")))
    env = {**os.environ, "SIMURG_PURE_PYTHON": "1" if a.pure_python else "0"}
    runs = []
    for t in (int(x) for x in a.threads.split(",")):
        out = a.out / f"t{t}"
        t0 = time.time()
        cmd = [a.simurg, "genotype", "--bpm", a.bpm, "--egt", a.egt, "--idats", str(idats), "--out", str(out),
               "--threads", str(t), "--no-cohort-qc", "--autocall-date", "1/1/2026 12:00:00 AM",
               "--quiet", "--force"]  # fmt: skip
        rc = subprocess.run(cmd, env=env).returncode
        wall = time.time() - t0
        setup, phase = phase_seconds(out / "logs" / "genotype.log")
        runs.append(
            {
                "threads": t,
                "rc": rc,
                "wall_s": round(wall, 1),
                "setup_s": round(setup, 1),
                "per_sample_phase_s": round(phase, 2),
                "samples_per_min": round(60 * n / phase, 1) if phase else None,
                "samples_per_min_per_core": round(60 * n / phase / t, 1) if phase else None,
            }
        )
        print(json.dumps(runs[-1]), flush=True)
    base = runs[0]["per_sample_phase_s"] * runs[0]["threads"]
    for r in runs:
        r["speedup_vs_first"] = round(base / r["per_sample_phase_s"], 2) if r["per_sample_phase_s"] else None
        r["efficiency"] = round(r["speedup_vs_first"] / r["threads"], 3) if r["speedup_vs_first"] else None
    from simurg._core import backend

    summary = {
        "samples": n,
        "backend": "python" if a.pure_python else backend(),
        "cpu_count": os.cpu_count(),
        "machine": os.uname().machine,
        "runs": runs,
    }
    (a.out / "scaling.json").write_text(json.dumps(summary, indent=1))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())