"""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 lightgbm  # noqa: F401  (before torch: libomp segfault on macOS)
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)
        t0 = time.time(); n = 0
        for i, b in enumerate(loader):
            n += b["y"].shape[0]
            if i >= 30:
                break
        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()