"""Engine to proxy features on the live capture code path, no browser anywhere.

Per chunk (12 frames, 3 condition frames), no pacing:
  chunk_to_features: engine command out -> features of all 3 condition frames ready
  last_frame_to_features: last condition frame's scene arrives -> its features ready
Also checks the features against the runner's way (one proxy pass over the 3
stacked frames) bit for bit.
"""
import json
import os
import statistics
import sys
import time
from pathlib import Path

import torch

EDEN = Path('/root/autodl-tmp/edenlive')
sys.path.insert(0, str(EDEN))
os.environ.setdefault('EDEN_GAMES_ROOT', str(EDEN / 'EdenLive/realtime/games'))
from eden.apps.game2v.game_timeline import GameTimeline
import native_capture

renderer = native_capture.NativeSceneRenderer('eiffel-assault', width=960, height=544)
MEASURE = len(sys.argv) > 1 and sys.argv[1] == 'frames'
renderer.measure_frames = MEASURE
reference = native_capture._proxy_encoder()
timeline = GameTimeline('eiffel-assault')
script = {1: [{'action': 'enter', 'phase': 'down'}, {'action': 'enter', 'phase': 'up'}],
          6: [{'action': 'w', 'phase': 'down'}], 14: [{'action': 'w', 'phase': 'up'}],
          18: [{'action': 'j', 'phase': 'down'}, {'action': 'j', 'phase': 'up'}]}
rows, identical, compared = [], 0, 0
for c in range(70):
    chunk = timeline.begin_chunk(9 if c == 0 else 12, 24)
    queue = list(script.get(c, []))
    def drain():
        out = list(queue)
        queue.clear()
        return out
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    packet = renderer.render_chunk(chunk, duration_s=0.0, drain_events=drain)
    t1 = time.perf_counter()
    features = native_capture.FEATURES.pop(packet.chunk_id)
    rows.append(dict(chunk_to_features=(t1 - t0) * 1e3, engine=renderer.last_engine_s * 1e3,
                     last_frame_to_features=renderer.last_frame_to_features_s * 1e3))
    if c >= 10 and c % 5 == 0:
        stacked = torch.stack([f.condition_frame for f in packet.frames if f.condition_frame is not None]).permute(0, 3, 1, 2).contiguous()
        expected = reference.infer(stacked).clone()
        torch.cuda.synchronize()
        identical += int(torch.equal(features, expected))
        compared += 1
renderer.close()
steady = rows[10:]


def stat(key):
    xs = sorted(r[key] for r in steady)
    return dict(median=round(statistics.median(xs), 2), p95=round(xs[int(len(xs) * 0.95)], 2), min=round(xs[0], 2))
if MEASURE:
    xs = sorted(x * 1e3 for x in renderer.frame_to_features_s[30:])
    print(json.dumps(dict(mode='per condition frame: scene from engine -> its features ready (synchronized)', n=len(xs),
                          median_ms=round(statistics.median(xs), 2), p95_ms=round(xs[int(len(xs) * .95)], 2),
                          min_ms=round(xs[0], 2), max_ms=round(xs[-1], 2))))
    sys.exit(0)
print(json.dumps(dict(chunks=len(steady), features=list(features.shape),
                      features_identical_to_runner_path=f'{identical}/{compared}',
                      chunk_to_features_ms=stat('chunk_to_features'), engine_12_frames_ms=stat('engine'),
                      last_frame_to_features_ms=stat('last_frame_to_features'))))