#!/usr/bin/env python

from __future__ import annotations

import os
import pathlib
import shlex
import subprocess
import tarfile

if os.getenv("SYSTEM") == "spaces":
    subprocess.run(shlex.split("pip install click==7.1.2"))
    subprocess.run(shlex.split("pip install typer==0.9.4"))

    import mim

    mim.uninstall("mmcv-full", confirm_yes=True)
    mim.install("mmcv-full==1.5.0", is_yes=True)

    subprocess.call(shlex.split("pip uninstall -y opencv-python"))
    subprocess.call(shlex.split("pip uninstall -y opencv-python-headless"))
    subprocess.call(shlex.split("pip install opencv-python-headless==4.8.0.74"))


import gradio as gr

from model import AppModel

DESCRIPTION = """# [ViTPose](https://github.com/ViTAE-Transformer/ViTPose)

Related app: [https://huggingface.co/spaces/Gradio-Blocks/ViTPose](https://huggingface.co/spaces/Gradio-Blocks/ViTPose)
"""


def extract_tar() -> None:
    if pathlib.Path("mmdet_configs/configs").exists():
        return
    with tarfile.open("mmdet_configs/configs.tar") as f:
        f.extractall("mmdet_configs")


extract_tar()

model = AppModel()

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column():
            input_video = gr.Video(label="Input Video", format="mp4", elem_id="input_video")
            detector_name = gr.Dropdown(
                label="Detector", choices=list(model.det_model.MODEL_DICT.keys()), value=model.det_model.model_name
            )
            pose_model_name = gr.Dropdown(
                label="Pose Model", choices=list(model.pose_model.MODEL_DICT.keys()), value=model.pose_model.model_name
            )
            det_score_threshold = gr.Slider(label="Box Score Threshold", minimum=0, maximum=1, step=0.05, value=0.5)
            max_num_frames = gr.Slider(label="Maximum Number of Frames", minimum=1, maximum=300, step=1, value=60)
            predict_button = gr.Button("Predict")
            pose_preds = gr.State()

            paths = sorted(pathlib.Path("videos").rglob("*.mp4"))
            gr.Examples(examples=[[path.as_posix()] for path in paths], inputs=input_video)

        with gr.Column():
            result = gr.Video(label="Result", format="mp4", elem_id="result")
            vis_kpt_score_threshold = gr.Slider(
                label="Visualization Score Threshold", minimum=0, maximum=1, step=0.05, value=0.3
            )
            vis_dot_radius = gr.Slider(label="Dot Radius", minimum=1, maximum=10, step=1, value=4)
            vis_line_thickness = gr.Slider(label="Line Thickness", minimum=1, maximum=10, step=1, value=2)
            redraw_button = gr.Button("Redraw")

    detector_name.change(fn=model.det_model.set_model, inputs=detector_name)
    pose_model_name.change(fn=model.pose_model.set_model, inputs=pose_model_name)
    predict_button.click(
        fn=model.run,
        inputs=[
            input_video,
            detector_name,
            pose_model_name,
            det_score_threshold,
            max_num_frames,
            vis_kpt_score_threshold,
            vis_dot_radius,
            vis_line_thickness,
        ],
        outputs=[
            result,
            pose_preds,
        ],
    )
    redraw_button.click(
        fn=model.visualize_pose_results,
        inputs=[
            input_video,
            pose_preds,
            vis_kpt_score_threshold,
            vis_dot_radius,
            vis_line_thickness,
        ],
        outputs=result,
    )

if __name__ == "__main__":
    demo.queue(max_size=10).launch()