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 (type(t[0])(v) if isinstance(t, tuple) else v))
                          for v, (k, v0, t) in zip(args,
                                                   (y for _, x in CONFIG_SPEC for y in x))]))
    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(), '%Y-%m-%d-%H-%M-%S')}.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():
    with gr.Blocks() as ui:
        gr.Markdown("# Pulsar+CLIP")
        gr.Markdown("Generate 3D point clouds from text!")
        
        with gr.Group():
            gr.Markdown("## Settings")
            inputs = []
            defaults = []
            with gr.Tabs():
                for name, section in CONFIG_SPEC:
                    with gr.TabItem(name):
                        for k, v0, t in section:
                            if t in (float, int):
                                element = gr.Number(label=k, value=v0)
                            elif t == str:
                                element = gr.Textbox(label=k, value=v0)
                            elif t == bool:
                                element = gr.Checkbox(label=k, value=v0)
                            elif isinstance(t, tuple):
                                element = gr.Slider(*t, label=k, value=v0)
                            elif isinstance(t, list):
                                element = gr.Dropdown(label=k, value=v0, choices=t)
                            else:
                                raise TypeError(f"Input format {t} should be one of str, int, bool, tuple, list")
                                element = 1/0
                            inputs.append(element)
                            defaults.append(v0)
        
        with gr.Row():
            with gr.Column():
                button = gr.Button("Run")
                gr.Markdown("## Result")
                output = gr.Video()
                
                button.click(fn=generate, inputs=inputs, outputs=output)
        
            with gr.Column():
                gr.Markdown("## Examples")
                gr.Examples(fn=generate, inputs=inputs, outputs=output,
                            examples=[defaults])
    
    ui.launch()
    
    #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()