"""P4 profiling: data-loader throughput, training step throughput and peak memory per model preset and batch size,
inference latency per origin. Writes results/metrics/benchmark_<node>.json. Uses real cached grids (private)."""
from __future__ import annotations

import argparse
import json
import platform
import subprocess
import sys
import time
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from chronos_x.config import load_config  # noqa: E402
from chronos_x import pipeline as P  # noqa: E402
from chronos_x.data.windows import WindowConfig, OriginIndex  # noqa: E402
from chronos_x.models.model import MRPT, MRPTConfig  # noqa: E402
from chronos_x.train.sampling import OriginDataset, collate  # noqa: E402
from chronos_x.train.trainer import pick_device, to_device, mem_mb  # noqa: E402
from chronos_x.train.losses import forecast_loss  # noqa: E402
from chronos_x import experiment as E  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]


def main() -> None:
    import resource
    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE); resource.setrlimit(resource.RLIMIT_NOFILE, (min(hard, 65536), hard))
    torch.multiprocessing.set_sharing_strategy("file_system")
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/pilot.yaml")
    ap.add_argument("--symbols-from", default="pilot")
    ap.add_argument("--presets", default="small,spec,base")
    ap.add_argument("--batch-sizes", default="32,64,128")
    ap.add_argument("--steps", type=int, default=20)
    a = ap.parse_args()
    cfg = load_config(a.config)
    man = json.loads((ROOT / cfg["universe"]["manifest"]).read_text())
    rows = E.symbols_for(man, a.symbols_from)
    d = cfg["data"]
    dirs = P.grid_dirs(Path(cfg["paths"]["cache"]), d["start"], d["end"], [r["symbol"] for r in rows])
    wcfg = WindowConfig(**cfg["windows"])
    idx = OriginIndex(dirs, wcfg, None, stride=1, max_per_session=10, seed=0)
    ds = OriginDataset(idx)
    dev = pick_device()
    out = {"node": platform.node(), "chip": subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True).stdout.strip(),
           "torch": torch.__version__, "device": str(dev), "n_origins": len(idx), "loader": {}, "train": {}, "inference": {}}
    # data loader throughput (CPU workers)
    for nw in (0, 4, 8):
        loader = DataLoader(ds, batch_size=64, shuffle=True, num_workers=nw, collate_fn=collate, persistent_workers=nw > 0,
                            prefetch_factor=4 if nw > 0 else None)
        it = iter(loader)
        for _ in range(4 if nw else 1):           # warm-up: worker spawn, torch import, grid mmap
            next(it)
        t0 = time.time(); n = 0
        for i in range(40):
            b = next(it); n += b["y"].shape[0]
        out["loader"][f"workers_{nw}"] = {"ex_per_s": n / (time.time() - t0)}
        print("loader", nw, out["loader"][f"workers_{nw}"], flush=True)
    # training steps
    for preset in a.presets.split(","):
        for bs in (int(x) for x in a.batch_sizes.split(",")):
            model = MRPT(MRPTConfig.preset(preset)).to(dev)
            opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
            loader = DataLoader(ds, batch_size=bs, shuffle=True, num_workers=8, collate_fn=collate, drop_last=True)
            it = iter(loader); batches = [to_device(next(it), dev) for _ in range(6)]
            # warm-up
            for b in batches[:2]:
                loss, _ = forecast_loss(model(b), b, "student_t"); loss.backward(); opt.step(); opt.zero_grad()
            if dev.type == "mps":
                torch.mps.synchronize()
            t0 = time.time()
            for i in range(a.steps):
                b = batches[2 + i % 4]
                loss, _ = forecast_loss(model(b), b, "student_t"); loss.backward(); opt.step(); opt.zero_grad()
            if dev.type == "mps":
                torch.mps.synchronize()
            el = time.time() - t0
            rec = {"params": model.n_params(), "steps_per_s": a.steps / el, "ex_per_s": a.steps * bs / el, "peak_mem_mb": mem_mb(dev),
                   "loss_finite": bool(torch.isfinite(loss))}
            out["train"][f"{preset}_bs{bs}"] = rec
            print("train", preset, bs, rec, flush=True)
            # inference latency (batch 1 and batch 256)
            model.eval()
            with torch.no_grad():
                b1 = {k: v[:1] for k, v in batches[0].items()}
                for _ in range(3):
                    model(b1)
                if dev.type == "mps":
                    torch.mps.synchronize()
                t0 = time.time()
                for _ in range(20):
                    model(b1)
                if dev.type == "mps":
                    torch.mps.synchronize()
                lat1 = (time.time() - t0) / 20
                t0 = time.time()
                for b in batches[:4]:
                    model(b)
                if dev.type == "mps":
                    torch.mps.synchronize()
                latb = (time.time() - t0) / (4 * bs)
            out["inference"][f"{preset}_bs{bs}"] = {"latency_ms_batch1": lat1 * 1e3, "latency_ms_per_origin_batched": latb * 1e3}
            del model, opt
            if dev.type == "mps":
                torch.mps.empty_cache()
    (ROOT / "results" / "metrics" / f"benchmark_{platform.node().split('.')[0]}.json").write_text(json.dumps(out, indent=1))
    print(json.dumps(out["train"], indent=1))


if __name__ == "__main__":
    main()