"""Can the proxy encoder run one condition frame at a time without changing the features?

Encodes the same 3 rendered condition frames as one batch (what generate() does
today) and one frame at a time, compares the outputs bit for bit, and times both
with CUDA events, eager and with the encoder's CUDA graph.
"""
import json
import os
import sys
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
from lightx2v.models.input_encoders.edenlive.proxy_dino import ProxyDINO
import native_capture

profile = json.loads((EDEN / 'runtime-5090.json').read_text())
renderer = native_capture.NativeSceneRenderer('eiffel-assault', width=960, height=544)
timeline = GameTimeline('eiffel-assault')
batches = []
for c in range(12):
    start = [{'action': 'enter', 'phase': 'down'}, {'action': 'enter', 'phase': 'up'}] if c == 1 else []
    packet = renderer.render_chunk(timeline.begin_chunk(9 if c == 0 else 12, 24), drain_events=lambda: start)
    batches.append(torch.stack([f.condition_frame for f in packet.frames if f.condition_frame is not None])
                   .permute(0, 3, 1, 2).contiguous())
renderer.close()


def timed(fn, repeat=20):
    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
    for _ in range(3):
        fn()
    torch.cuda.synchronize()
    start.record()
    for _ in range(repeat):
        fn()
    end.record()
    end.synchronize()
    return start.elapsed_time(end) / repeat


report = {}
for graph in (False, True):
    proxy = ProxyDINO(str(EDEN / 'models/EdenLive-x2v'), profile.get('proxy_resize_mode', 'bilinear'),
                      use_cuda_graph=graph, attention_type='sage_attn2', quantized_ckpt=profile.get('proxy_quantized_ckpt'))
    identical, max_diff = 0, 0.0
    for x in batches[4:]:
        whole = proxy.infer(x).clone()
        split = torch.cat([proxy.infer(x[i:i + 1]).clone() for i in range(3)], dim=2)
        diff = (whole.float() - split.float()).abs().max().item()
        identical += int(torch.equal(whole, split))
        max_diff = max(max_diff, diff)
    x = batches[-1]
    report['cuda_graph' if graph else 'eager'] = dict(
        identical_chunks=f'{identical}/{len(batches) - 4}', max_abs_diff=max_diff,
        batch3_ms=round(timed(lambda: proxy.infer(x)), 2),
        one_frame_ms=round(timed(lambda: proxy.infer(x[:1])), 2),
        output=list(proxy.infer(x).shape))
    del proxy
    torch.cuda.empty_cache()
print(json.dumps(report))