import numpy as np
import PIL.Image as Image
import os
import time
import torch
from flowsrepo import *


examples = {
    "satellite": SatelliteFlow,
    "dragons": BattleOfDragonsFlow,
    "earth": SphereFlow,
    "meltingman": MeltingManFlow,
    "glass": GlassFlow,
}

timestamp = time.time()
N = 512

for examplename in examples:
    image_warper = examples[examplename](N=N)
    image = image_warper.get_default_image()
    image = image.resize((N, N))
    image = torch.from_numpy(np.array(image) / 255).float().permute(2, 0, 1)
    framesteps = image_warper.get_default_framesteps()

    frames = image_warper.warp(
        should_apply_to_last_frame=True,
        image = image,
        align_corners=True,
    )

    npframes = (frames * 255).byte().permute(0, 2, 3, 1).cpu().numpy()
    dirpath = f"output/pixelspace/{timestamp}/{time.time()}_{examplename}"
    os.makedirs(dirpath, exist_ok=True)
    for i, frame in enumerate(npframes):
        Image.fromarray(frame).save(f"{dirpath}/frame_{i:03d}.png")
