#!/usr/bin/env python

from __future__ import annotations

import argparse

import gradio as gr

from model import Model

TITLE = '# Anime Face Generation with [Diffusers](https://github.com/huggingface/diffusers)'
DESCRIPTION = 'Expected execution time on Hugging Face Spaces: 5s (DDIM, 20 steps), 6s (PNDM, 20 steps), 247s (DDPM, 1000 steps)'
FOOTER = '<img id="visitor-badge" src="https://visitor-badge.glitch.me/badge?page_id=hysts.diffusers-anime-faces" alt="visitor badge" />'


def get_sample_image_url(file_name: str) -> str:
    sample_image_dir = 'https://huggingface.co/spaces/hysts/diffusers-anime-faces/resolve/main/samples'
    return f'{sample_image_dir}/{file_name}'


def get_sample_image_markdown(name: str) -> str:
    model_name = name.split()[0]
    if name == 'ddpm-128-exp000 (DDPM)':
        scheduler = 'DDPM'
        steps = 1000
        file_name = f'{model_name}.png'
    elif name == 'ddpm-128-exp000 (DDIM, 20 steps)':
        scheduler = 'DDIM'
        steps = 20
        file_name = f'{model_name}-ddim-20steps.png'
    else:
        raise ValueError
    url = get_sample_image_url(file_name)
    text = f'''
            - size: 128x128
            - seed: 0-99
            - scheduler: {scheduler}
            - steps: {steps}

            ![sample images]({url})'''
    return text


def update_scheduler_type(name: str) -> dict:
    visible = name != 'DDPM'
    if name == 'PNDM':
        minimum = 4
        maximum = 100
    else:
        minimum = 1
        maximum = 200
    return gr.Slider.update(visible=visible,
                            minimum=minimum,
                            maximum=maximum,
                            value=20)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', type=str, default='cpu')
    args = parser.parse_args()

    model = Model(args.device)

    with gr.Blocks(css='style.css') as demo:
        gr.Markdown(TITLE)

        with gr.Tabs():
            with gr.TabItem('Simple Mode'):
                run_button_simple = gr.Button('Generate')
                result_simple = gr.Image(show_label=False,
                                         elem_id='result-grid')

            with gr.TabItem('Advanced Mode'):
                gr.Markdown(DESCRIPTION)

                with gr.Row():
                    with gr.Column():
                        with gr.Group():
                            model_name = gr.Dropdown(
                                model.MODEL_NAMES,
                                value=model.MODEL_NAMES[0],
                                label='Model',
                                interactive=False)
                            scheduler_type = gr.Radio(
                                choices=['DDPM', 'DDIM', 'PNDM'],
                                value='DDIM',
                                label='Scheduler')
                            num_steps = gr.Slider(1,
                                                  200,
                                                  step=1,
                                                  value=20,
                                                  label='Number of Steps')
                            seed = gr.Slider(0,
                                             100000,
                                             step=1,
                                             value=1234,
                                             label='Seed')
                            run_button = gr.Button('Run')
                    with gr.Column():
                        result = gr.Image(show_label=False, elem_id='result')

            with gr.TabItem('Sample Images'):
                with gr.Row():
                    model_name2 = gr.Dropdown([
                        'ddpm-128-exp000 (DDPM)',
                        'ddpm-128-exp000 (DDIM, 20 steps)',
                    ],
                                              value='ddpm-128-exp000 (DDPM)',
                                              label='Model')
                with gr.Row():
                    text = get_sample_image_markdown(model_name2.value)
                    sample_images = gr.Markdown(text)

        gr.Markdown(FOOTER)

        model_name.change(fn=model.set_pipeline,
                          inputs=[
                              model_name,
                              scheduler_type,
                          ],
                          outputs=None)
        scheduler_type.change(fn=update_scheduler_type,
                              inputs=scheduler_type,
                              outputs=num_steps,
                              queue=False)
        scheduler_type.change(fn=model.set_pipeline,
                              inputs=[
                                  model_name,
                                  scheduler_type,
                              ],
                              outputs=None)
        run_button_simple.click(fn=model.run_simple,
                                inputs=None,
                                outputs=result_simple)
        run_button.click(fn=model.run,
                         inputs=[
                             model_name,
                             scheduler_type,
                             num_steps,
                             seed,
                         ],
                         outputs=result)
        model_name2.change(fn=get_sample_image_markdown,
                           inputs=model_name2,
                           outputs=sample_images)

    demo.launch(enable_queue=True, share=False)


if __name__ == '__main__':
    main()