"""Streaming VAE timings at 960x544: encoder (1 frame, then 4-frame chunks) and decoder (1 latent -> 4 frames)."""
import json
import time

import torch

from sftrain import conditions

device = torch.device('cuda')
enc = conditions.vae_encoder(device)
dec = conditions.vae_decoder(device)
H, W = conditions.HEIGHT, conditions.WIDTH
print(json.dumps(dict(H=H, W=W, enc_dtype=str(getattr(enc.config, 'dtype', '?')), dec_dtype=str(dec.config.dtype))), flush=True)


def timed(fn, n):
    out = []
    for _ in range(n):
        torch.cuda.synchronize()
        t = time.perf_counter()
        fn()
        torch.cuda.synchronize()
        out.append((time.perf_counter() - t) * 1000)
    return out


with torch.no_grad():
    for rep in range(2):
        cache = enc.initialize_autoregressive_cache()
        first = torch.rand(1, 1, 3, H, W, device=device) * 2 - 1
        t0 = timed(lambda: enc(first, autoregressive_index=0, cache=cache), 1)
        ts = []
        for i in range(1, 11):
            chunk = torch.rand(1, 4, 3, H, W, device=device) * 2 - 1
            ts += timed(lambda: enc(chunk, autoregressive_index=i, cache=cache), 1)
        full = torch.rand(1, 81, 3, H, W, device=device) * 2 - 1
        tf = timed(lambda: enc(full, autoregressive_index=0, cache=enc.initialize_autoregressive_cache()), 1)
        lat = torch.randn(1, 11, 48, conditions.LATENT_H, conditions.LATENT_W, device=device, dtype=dec.config.dtype)
        dcache = dec.initialize_autoregressive_cache()
        d0 = timed(lambda: dec(lat[:, :1], autoregressive_index=0, cache=dcache), 1)
        ds = []
        for i in range(1, 11):
            ds += timed(lambda: dec(lat[:, i:i + 1], autoregressive_index=i, cache=dcache), 1)
        print(json.dumps(dict(rep=rep, enc_first_ms=round(t0[0], 1), enc_4frame_ms=[round(x, 1) for x in ts],
                              enc_81frames_ms=round(tf[0], 1), dec_first_ms=round(d0[0], 1),
                              dec_1latent_4frames_ms=[round(x, 1) for x in ds])), flush=True)