#!/usr/bin/env python3
"""Instrumented 100 ms canonical-backend cost audit."""
from __future__ import annotations

import json
import os
import resource
import subprocess
import threading
import time
from pathlib import Path

import numpy as np
import torch
import yaml

import phase4p


def main() -> None:
    model = Path("/data/digital-drosophila/models/male-cns-lif/deterministic-runtime-v2")
    cfg = yaml.safe_load((model / "phase4p.yaml").read_text()); exp = Path(cfg["paths"]["experiment_root"])
    rt = phase4p.CanonicalRuntime(cfg, model, exp); spec = cfg["reference_trials"][1]
    target_np = np.flatnonzero(rt.pops[spec["population"]]).astype(np.int64)
    targets = torch.as_tensor(target_np, dtype=torch.int64, device=rt.device); steps = 1000
    forced = phase4p.phase4c.generate_forced(steps, len(target_np), float(spec["rate_hz"]), rt.dt, int(spec["seed"]), rt.device)
    rt.reset(); samples = []; stop = threading.Event()
    def sample():
        while not stop.is_set():
            try:
                text = subprocess.check_output(["nvidia-smi", "--id=" + cfg["device"]["visible_gpu_uuid"],
                    "--query-gpu=utilization.gpu,memory.used", "--format=csv,noheader,nounits"], text=True).strip()
                u, m = text.split(","); samples.append({"time": time.time(), "gpu_util_percent": float(u), "gpu_memory_mib": float(m)})
            except Exception as e: samples.append({"error": str(e)})
            stop.wait(0.10)
    thread = threading.Thread(target=sample, daemon=True); thread.start()
    use0 = resource.getrusage(resource.RUSAGE_SELF); cpu0 = use0.ru_utime + use0.ru_stime
    torch.cuda.reset_peak_memory_stats(); torch.cuda.synchronize()
    wall0 = time.perf_counter(); state_submit = force_submit = sort_submit = reduce_submit = 0.0; contributions = 0
    for step in range(steps):
        slot = step % rt.qlen; t = time.perf_counter()
        rt.ext.state_step(rt.v, rt.g, rt.ref, rt.queue[slot], rt.spike_u8, rt.counts,
                          float(rt.gd), float(rt.alpha), float(rt.rest), float(rt.resetv), float(rt.threshold), rt.ref_steps)
        state_submit += time.perf_counter() - t; t = time.perf_counter()
        rt.ext.force_step(rt.v, rt.ref, rt.spike_u8, rt.counts, targets, forced[step], float(rt.resetv), rt.ref_steps)
        force_submit += time.perf_counter() - t
        sources = torch.nonzero(rt.spike_u8, as_tuple=False).flatten()
        st = rt.canonical_fanout(sources, rt.queue[(step + rt.delay_steps) % rt.qlen])
        sort_submit += st["sort_seconds"]; reduce_submit += st["reduce_seconds"]; contributions += st["contributions"]
    torch.cuda.synchronize(); wall = time.perf_counter() - wall0
    use1 = resource.getrusage(resource.RUSAGE_SELF); cpu1 = use1.ru_utime + use1.ru_stime
    stop.set(); thread.join()
    good = [x for x in samples if "gpu_util_percent" in x]
    out = {"trial": "auditory-100hz", "biological_ms": 100.0, "steps": steps, "wall_seconds": wall,
           "wall_per_100ms_biological": wall, "realtime_factor": wall / 0.1, "steps_per_second": steps / wall,
           "process_CPU_percent_of_one_core": 100 * (cpu1 - cpu0) / wall,
           "GPU_utilization_percent_mean": float(np.mean([x["gpu_util_percent"] for x in good])) if good else None,
           "GPU_utilization_percent_max": float(np.max([x["gpu_util_percent"] for x in good])) if good else None,
           "nvidia_smi_memory_mib_max": float(np.max([x["gpu_memory_mib"] for x in good])) if good else None,
           "torch_peak_VRAM_bytes": int(torch.cuda.max_memory_allocated()), "contribution_count": int(contributions),
           "CPU_submit_seconds": {"state_integration": state_submit, "external_force": force_submit,
                                  "canonical_sort_and_segment_setup": sort_submit, "segmented_reduction": reduce_submit},
           "sampling_interval_seconds": 0.10, "samples": samples,
           "interpretation": "Submission sub-times overlap asynchronous CUDA work; synchronized total wall time is authoritative."}
    p = exp / "performance/instrumented-100ms.json"; p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(out, indent=2, sort_keys=True) + "\n"); print(json.dumps(out, indent=2), flush=True)


if __name__ == "__main__": main()