#!/usr/bin/env python3
"""Per-step timing of Ming-Image's DiT with allocator stats and a GPU clock/power sampler.

Diagnoses step time that grows within one generation. For every DiT call it records the
synchronized wall time, the caching allocator's reserved/allocated bytes, how many device
mallocs and malloc retries (fragmentation) have happened so far; a sampler thread reads the
GPU sclk, power and temperature twice a second.

  usage: PYTHONPATH=<code_dir> step_probe.py --prompt P.json [--runs N] -- <infer.py args>
"""
import argparse
import glob
import json
import sys
import threading
import time


def read_gpu():
    base = "/sys/class/drm/card0/device"
    sclk = next((l.split(":")[1].strip().rstrip("*").strip() for l in open(f"{base}/pp_dpm_sclk") if "*" in l), "?")
    hw = sorted(glob.glob(f"{base}/hwmon/hwmon*"))[0]
    power = int(open(f"{hw}/power1_average").read()) / 1e6
    temp = int(open(f"{hw}/temp1_input").read()) / 1e3
    busy = int(open(f"{base}/gpu_busy_percent").read())
    return sclk, power, temp, busy


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompt", required=True)
    ap.add_argument("--runs", type=int, default=1)
    own, rest = ap.parse_known_args()
    if rest and rest[0] == "--":
        rest = rest[1:]
    sys.argv = [sys.argv[0], "--prompt", own.prompt] + rest

    import torch
    import infer

    args = infer.parse_args()
    model_directory = infer.resolve_model_directory(args.model, local_files_only=True)
    caps = infer.load_checkpoint_capabilities(model_directory)
    resolution = infer.resolve_task_resolution(args.task, args.resolution)
    sampling = caps.resolve_sampling_parameters(steps=args.steps, cfg=args.cfg)
    dtype = infer._dtype(args.dtype)
    model, processor = infer.load_model_and_processor(model_directory, args)
    prompt = infer._load_prompt(own.prompt)

    samples, stop = [], threading.Event()

    def sampler():
        t0 = time.perf_counter()
        while not stop.is_set():
            samples.append((round(time.perf_counter() - t0, 1),) + read_gpu())
            time.sleep(0.5)

    dit = model.diffusion_loss.train_model
    marks = {}

    def pre(_module, _args, _kwargs):
        torch.cuda.synchronize()
        marks["t"] = time.perf_counter()

    def post(_module, _args, _kwargs, _out):
        torch.cuda.synchronize()
        st = torch.cuda.memory_stats()
        step_log.append({
            "step_s": round(time.perf_counter() - marks["t"], 2),
            "reserved_gib": round(torch.cuda.memory_reserved() / 2**30, 2),
            "allocated_gib": round(torch.cuda.memory_allocated() / 2**30, 2),
            "device_mallocs": st.get("num_device_alloc", 0),
            "device_frees": st.get("num_device_free", 0),
            "alloc_retries": st.get("num_alloc_retries", 0),
        })

    dit.register_forward_pre_hook(pre, with_kwargs=True)
    dit.register_forward_hook(post, with_kwargs=True)
    thread = threading.Thread(target=sampler, daemon=True)
    thread.start()
    for run in range(own.runs):
        step_log = []
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        infer.run_generation(model, processor, caps, task=args.task, prompt=prompt, input_image=None,
                             resolution=resolution, sampling=sampling, seed=args.seed,
                             num_layers=args.num_layers, dtype=dtype)
        torch.cuda.synchronize()
        print(f"RUN {run} total_s {time.perf_counter() - t0:.1f}", flush=True)
        for i, row in enumerate(step_log):
            print("STEP " + json.dumps({"run": run, "i": i, **row}), flush=True)
    stop.set()
    thread.join()
    for s in samples[:: max(1, len(samples) // 60)]:
        print("GPU t=%6.1fs sclk=%s power=%.0fW temp=%.0fC busy=%d%%" % s)


if __name__ == "__main__":
    main()