#!/usr/bin/env python3
# Published copy. Host names and paths are replaced by <placeholders> or left out of comments, nothing else is changed. The sha256 of the original is in sha256.txt.
"""Idol timing on one GPU (RTX 5060 Ti 16 GB or RTX 5090 32 GB), author defaults (50 steps, CFG 3.0, seed 42, bf16 DINOv3, fast sampling on).
Run: CUDA_VISIBLE_DEVICES=<n> IDOL_PHYS_GPU=<n> HF_HUB_OFFLINE=1 HF_HUB_CACHE=<hf-cache> python bench_idol.py <out_dir> <image_dir>...
First image of the process (torch.compile) is timed separately, the rest are warm runs (median).
Settings: 262,144 Gaussians (author default), then 1,048,576 (author maximum) for the peak VRAM number.
External peak = nvidia-smi memory.used of the physical GPU IDOL_PHYS_GPU sampled every 0.2 s (includes the CUDA context); torch peak = max_memory_allocated."""
import json, os, statistics, subprocess, sys, threading, time

import torch

sys.path.insert(0, "<idol-sdk>")
from idol.pipeline import IdolPipeline  # noqa: E402

OUT = sys.argv[1]
DIRS = sys.argv[2:]
PHYS_GPU = os.environ.get("IDOL_PHYS_GPU", "1")
WEIGHTS = "<hf-cache>/models--speridlabs--idol/snapshots/2731c2291fc65a797c5cb627a5e5526fc1490902"
EXT = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
os.makedirs(OUT, exist_ok=True)

peak = {"mib": 0, "stop": False}


def poll():
    while not peak["stop"]:
        try:
            v = int(subprocess.run(["nvidia-smi", "-i", PHYS_GPU, "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
                                   capture_output=True, text=True, timeout=5).stdout.strip())
            peak["mib"] = max(peak["mib"], v)
        except Exception:
            pass
        time.sleep(0.2)


def smi_now():
    return int(subprocess.run(["nvidia-smi", "-i", PHYS_GPU, "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
                              capture_output=True, text=True).stdout.strip())


images = []
for d in DIRS:
    images += sorted(os.path.join(d, n) for n in os.listdir(d) if n.lower().endswith(EXT))
print("images:", len(images), flush=True)
res = {"started": time.strftime("%Y-%m-%d %H:%M:%S"), "gpu_name": torch.cuda.get_device_name(0), "torch": torch.__version__,
       "n_images": len(images), "idle_mib_before": smi_now()}
threading.Thread(target=poll, daemon=True).start()

t = time.time()
pipe = IdolPipeline.from_pretrained(WEIGHTS)
torch.cuda.synchronize()
res["load_s"] = round(time.time() - t, 2)
res["after_load_mib_smi"] = smi_now()
res["after_load_torch_alloc_mib"] = round(torch.cuda.memory_allocated() / 2**20)
print("loaded", res["load_s"], "s", flush=True)

for setting, ng in (("g262144", 262144), ("g1048576", 1048576)):
    odir = os.path.join(OUT, setting)
    os.makedirs(odir, exist_ok=True)
    torch.cuda.reset_peak_memory_stats()
    peak["mib"] = 0
    times = []
    todo = images if ng == 262144 else images[::max(1, len(images) // int(os.environ.get("IDOL_N1M", "8")))][:int(os.environ.get("IDOL_N1M", "8"))]
    for i, p in enumerate(todo):
        name = os.path.splitext(os.path.basename(p))[0]
        torch.cuda.synchronize()
        t = time.time()
        out = pipe(p, mode="synthetic", seed=42, steps=50, cfg_strength=3.0, rescale_t=3.0, num_gaussians=ng)
        torch.cuda.synchronize()
        dt = time.time() - t
        out.gaussian.save_ply(os.path.join(odir, name + ".ply"))
        out.preprocessed_image.save(os.path.join(odir, name + "_input.png"))
        with open(os.path.join(odir, name + "_camera.json"), "w") as f:
            json.dump(out.camera, f, indent=2)
        times.append((name, round(dt, 3)))
        print(setting, i, name, round(dt, 3), "s", flush=True)
    cold = times[0][1] if setting == "g262144" else None
    warm = [x[1] for x in (times[1:] if setting == "g262144" else times)]
    res[setting] = {"num_gaussians": ng, "n_images": len(todo), "times": times, "first_image_s_cold_compile": cold,
                    "warm_median_s": round(statistics.median(warm), 3), "warm_min_s": min(warm), "warm_max_s": max(warm), "warm_n": len(warm),
                    "torch_peak_alloc_mib": round(torch.cuda.max_memory_allocated() / 2**20),
                    "torch_peak_reserved_mib": round(torch.cuda.max_memory_reserved() / 2**20),
                    "smi_peak_mib": peak["mib"]}
    json.dump(res, open(os.path.join(OUT, "results.json"), "w"), indent=1)

peak["stop"] = True
res["finished"] = time.strftime("%Y-%m-%d %H:%M:%S")
json.dump(res, open(os.path.join(OUT, "results.json"), "w"), indent=1)
print("DONE", flush=True)