# coding: utf-8

"""
The entrance of the gradio
"""

import tyro
import gradio as gr
import spaces
import os.path as osp
from src.utils.helper import load_description
from src.gradio_pipeline import GradioPipeline
from src.config.crop_config import CropConfig
from src.config.argument_config import ArgumentConfig
from src.config.inference_config import InferenceConfig


def partial_fields(target_class, kwargs):
    return target_class(**{k: v for k, v in kwargs.items() if hasattr(target_class, k)})


# set tyro theme
tyro.extras.set_accent_color("bright_cyan")
args = tyro.cli(ArgumentConfig)

# specify configs for inference
inference_cfg = partial_fields(InferenceConfig, args.__dict__)  # use attribute of args to initial InferenceConfig
crop_cfg = partial_fields(CropConfig, args.__dict__)  # use attribute of args to initial CropConfig

gradio_pipeline = GradioPipeline(
    inference_cfg=inference_cfg,
    crop_cfg=crop_cfg,
    args=args
)

@spaces.GPU()
def gpu_wrapped_execute_s_video(*args, **kwargs):
    return gradio_pipeline.execute_s_video(*args, **kwargs)

# assets
title_md = "assets/gradio_title.md"
example_portrait_dir = "assets/examples/source"
example_video_dir = "assets/examples/driving"

#################### interface logic ####################
css="""
#col-container {
    max_width: 1400px;
    margin: 0 auto;
}
"""
with gr.Blocks(css=css) as demo:

    with gr.Column(elem_id="col-container"):
        gr.HTML(load_description(title_md))

        gr.Markdown("""
        ## 🤗 This is the gradio demo for Vid2Vid LivePortrait.
        Please upload or use a webcam to get a Source Portrait Video (any aspect ratio) and upload a Driving Video (1:1 aspect ratio, or any aspect ratio with do crop (driving video) checked).
        """)

    with gr.Row():
        with gr.Column():
        
            source_video_input = gr.Video(label="Portrait Video Source")
            gr.Examples(
                examples=[
                    [osp.join(example_video_dir, "d0.mp4")],
                    [osp.join(example_video_dir, "d18.mp4")],
                    [osp.join(example_video_dir, "d19.mp4")],
                    [osp.join(example_video_dir, "d14.mp4")],
                    [osp.join(example_video_dir, "d6.mp4")],
                ],
                inputs=[source_video_input],
                cache_examples=False,
            )
            gr.Markdown("——————————")
            video_input = gr.Video(label="Driving Portrait Video")
            gr.Examples(
                examples=[
                    [osp.join(example_video_dir, "d0.mp4")],
                    [osp.join(example_video_dir, "d18.mp4")],
                    [osp.join(example_video_dir, "d19.mp4")],
                    [osp.join(example_video_dir, "d14.mp4")],
                    [osp.join(example_video_dir, "d6.mp4")],
                ],
                inputs=[video_input],
                cache_examples=False,
            )
            gr.Markdown(load_description("assets/gradio_description_animation.md"))
            with gr.Accordion(open=False, label="source Animation Instructions and Options"):
                gr.Markdown(load_description("assets/gradio_description_animation.md"))
                with gr.Row():
                    flag_relative_input = gr.Checkbox(value=True, label="relative motion")
                    flag_do_crop_input = gr.Checkbox(value=True, label="do crop (source)")
                    flag_remap_input = gr.Checkbox(value=True, label="paste-back")
                    flag_crop_driving_video_input = gr.Checkbox(value=False, label="do crop (driving video)")
            with gr.Row():
                process_button_source_animation = gr.Button("🚀 Animate video", variant="primary")
        with gr.Column():
            output_video1 = gr.Video(label="The animated video in the original image space")
            output_video_concat1 = gr.Video(label="The animated video")

    # binding functions for buttons
   
    process_button_source_animation.click(
        fn=gpu_wrapped_execute_s_video,
        inputs=[
            source_video_input,
            video_input,
            flag_relative_input,
            flag_do_crop_input,
            flag_remap_input,
            flag_crop_driving_video_input
        ],
        outputs=[output_video1, output_video_concat1],
        show_progress=True,
        show_api=False
    )


demo.queue(max_size=10).launch(show_api=False)