from pulsar_clip import PulsarCLIP, CONFIG_SPEC
from datetime import datetime
import gradio as gr


def generate(*args):
    pc = PulsarCLIP(dict([(k, t(v) if not isinstance(t, (tuple, list)) else v)
                          for v, (k, v0, t) in zip(args, CONFIG_SPEC)]))
    frames = []
    for image in pc.generate():
        frames.append(image)
    from tqdm.auto import tqdm
    from subprocess import Popen, PIPE
    fps = 30
    video_path = f"{datetime.strftime(datetime.now())}.mp4"
    if frames:
        p = Popen((f"ffmpeg -y -f image2pipe -vcodec png -r {fps} -i - -vcodec libx264 -r {fps} "
                   f"-pix_fmt yuv420p -crf 17 -preset fast ").split() + [str(video_path)], stdin=PIPE)
        for im in tqdm(frames):
            im.save(p.stdin, "PNG")
        p.stdin.close()
        p.wait()
    return video_path


def main():
    gr.Interface(inputs=[
        (gr.inputs.Number(label=k, default=v0) if t in (float, int) else
         gr.inputs.Checkbox(label=k, default=v0) if t == bool else gr.inputs.Textbox(label=k, default=v0) if t == str
         else gr.inputs.Dropdown(label=k, default=v0, choices=t) if isinstance(t, (tuple, list)) else 1/0)
        for k, v0, t in CONFIG_SPEC], outputs=gr.outputs.Video(), fn=generate).launch()


if __name__ == '__main__':
    main()