"""Time the browser-free path end to end: engine step + native GPU render, per condition frame.

usage: bench_render.py [frames] [compile 0|1]
GPU time comes from CUDA events around render(); wall time includes the engine
round trip and a synchronize, i.e. "frame ready for the proxy encoder".
"""
import json
import sys
import time

import torch

from engine_client import Engine
from native_render import NativeRenderer

FRAMES = int(sys.argv[1]) if len(sys.argv) > 1 else 480
COMPILE = (sys.argv[2] if len(sys.argv) > 2 else '1') == '1'

engine = Engine()
renderer = NativeRenderer(engine.scene, engine.scene_arrays, compile=COMPILE)
rows = []
for i in range(FRAMES):
    target = round((i + 1) * 60 / 24)
    events = [{'action': 'enter', 'phase': 'down'}, {'action': 'enter', 'phase': 'up'}] if i == 5 else []
    if 40 <= i < 440 and i % 24 == 0:
        events.append({'action': 'w', 'phase': 'down' if (i // 24) % 2 == 0 else 'up'})
    if 80 <= i < 440 and i % 37 == 0:
        events += [{'action': 'j', 'phase': 'down'}, {'action': 'j', 'phase': 'up'}]
    capture = i % 4 == 0
    t0 = time.perf_counter()
    header, frame = engine.advance(target, events, capture=capture, state=False)
    t1 = time.perf_counter()
    if not capture:
        rows.append(dict(capture=False, engine=(t1 - t0) * 1e3))
        continue
    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
    start.record()
    image = renderer.render(frame)
    end.record()
    t2 = time.perf_counter()
    end.synchronize()
    t3 = time.perf_counter()
    rows.append(dict(capture=True, engine=(t1 - t0) * 1e3, submit=(t2 - t1) * 1e3, ready=(t3 - t1) * 1e3,
                     gpu=start.elapsed_time(end), total=(t3 - t0) * 1e3))
engine.close()


def stat(xs):
    xs = sorted(xs)
    return dict(n=len(xs), median=round(xs[len(xs) // 2], 3), p95=round(xs[int(len(xs) * .95)], 3), max=round(xs[-1], 3))
caps = [r for r in rows if r['capture']][10:]   # first renders include compilation
sims = [r for r in rows if not r['capture']][30:]
print(json.dumps(dict(
    compile=COMPILE, image=list(image.shape), dtype=str(image.dtype), device=str(image.device),
    sim_frame_ms=stat([r['engine'] for r in sims]),
    condition_frame_engine_ms=stat([r['engine'] for r in caps]),
    render_submit_ms=stat([r['submit'] for r in caps]),
    render_gpu_ms=stat([r['gpu'] for r in caps]),
    render_ready_ms=stat([r['ready'] for r in caps]),
    condition_frame_total_ms=stat([r['total'] for r in caps]),
    gpu_memory_mb=round(torch.cuda.max_memory_allocated() / 2**20))))