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


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
    filename = datetime.strftime(datetime.now(), "%Y-%m-%d-%H-%M-%S")
    video_path = f"{filename}.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()
    model_path = f"{filename}.obj"
    pc.save_obj(model_path)
    # model_path = None  # TODO
    return [video_path, model_path, model_path]


def main():
    with gr.Blocks() as ui:
        gr.Markdown("# Pulsar+CLIP")
        gr.Markdown("<a target='_blank' href='https://colab.research.google.com/drive/1IvV3HGoNjRoyAKIX-aqSWa-t70PW3nPs'> <img src='https://colab.research.google.com/assets/colab-badge.svg' alt='Open In Colab'/> </a> [![arXiv](https://img.shields.io/badge/arXiv-2004.07484-b31b1b.svg)](https://arxiv.org/abs/2004.07484)")
        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)
        
        button = gr.Button("Run")
        gr.Markdown("## Result")
        with gr.Row():
            with gr.Column():
                video_result = gr.Video()
            with gr.Column():
                model_demo = gr.Model3D()
                model_file = gr.File()
        
        button.click(fn=generate, inputs=inputs, outputs=[video_result, model_demo, model_file])
        
        gr.Markdown("## Examples")
        gr.Examples(fn=generate, inputs=inputs, outputs=[video_result, model_demo, model_file],
                    examples=[defaults], cache_examples=True, examples_per_page=1)
    return ui

ui = main()
ui.configure_queue(concurrency_count=5).launch()
demo = ui