import os
os.system('python setup.py develop')

import argparse
from pathlib import Path

import gradio as gr
import matplotlib

from gradio_utils.utils import (process_img, get_select_coords, select_skeleton,
                                reset_skeleton, reset_kp, process)

LENGTH = 480  # Length of the square area displaying/editing images

matplotlib.use('agg')
model_dir = Path('./checkpoints')
parser = argparse.ArgumentParser(description='EdgeCape Demo')
parser.add_argument('--checkpoint',
                    help='checkpoint path',
                    default='ckpt/1shot_split1.pth')
args = parser.parse_args()
checkpoint_path = args.checkpoint
device = 'cuda'
TIMEOUT = 80

with gr.Blocks() as demo:
    gr.Markdown('''
    # We introduce EdgeCape, a novel framework that overcomes these limitations by predicting the graph's edge weights which optimizes localization. 
    To further leverage structural priors, we propose integrating Markovian Structural Bias, which modulates the self-attention interaction between nodes based on the number of hops between them. 
    We show that this improves the model’s ability to capture global spatial dependencies. 
    Evaluated on the MP-100 benchmark, which includes 100 categories and over 20K images, 
    EdgeCape achieves state-of-the-art results in the 1-shot setting and leads among similar-sized methods in the 5-shot setting, significantly improving keypoint localization accuracy.
    ### [Paper](https://arxiv.org/pdf/2411.16665) | [Project Page](https://orhir.github.io/edge_cape/) 
    ## Instructions
    1. Upload an image of the object you want to pose.
    2. Mark keypoints on the image.
    3. Mark limbs on the image.
    4. Upload an image of the object you want to pose to the query image (**bottom**).
    5. Click **Evaluate** to pose the query image.
    ''')

    global_state = gr.State({
        "images": {},
        "points": [],
        "skeleton": [],
        "prev_point": None,
        "curr_type_point": "start",
    })
    with gr.Row():
        # Upload & Preprocess Image Column
        with gr.Column():
            gr.Markdown(
                """<p style="text-align: center; font-size: 20px">Upload & Preprocess Image</p>"""
            )
            support_image = gr.Image(
                height=LENGTH,
                width=LENGTH,
                type="pil",
                image_mode="RGB",
                label="Preprocess Image",
                show_label=True,
                interactive=True,
            )

        # Click Points Column
        with gr.Column():
            gr.Markdown(
                """<p style="text-align: center; font-size: 20px">Click Points</p>"""
            )
            kp_support_image = gr.Image(
                type="pil",
                label="Keypoints Image",
                show_label=True,
                height=LENGTH,
                width=LENGTH,
                interactive=False,
                show_fullscreen_button=False,
            )
            with gr.Row():
                confirm_kp_button = gr.Button("Confirm Clicked Points", scale=3)
            with gr.Row():
                undo_kp_button = gr.Button("Undo Clicked Points", scale=3)

        # Editing Results Column
        with gr.Column():
            gr.Markdown(
                """<p style="text-align: center; font-size: 20px">Click Skeleton</p>"""
            )
            skel_support_image = gr.Image(
                type="pil",
                label="Skeleton Image",
                show_label=True,
                height=LENGTH,
                width=LENGTH,
                interactive=False,
                show_fullscreen_button=False,
            )
            with gr.Row():
                pass
            with gr.Row():
                undo_skel_button = gr.Button("Undo Skeleton")

    with gr.Row():
        with gr.Column():
            gr.Markdown(
                """<p style="text-align: center; font-size: 20px">Query Image</p>"""
            )
            query_image = gr.Image(
                type="pil",
                image_mode="RGB",
                label="Query Image",
                show_label=True,
                interactive=True,
            )
        with gr.Column():
            gr.Markdown(
                """<p style="text-align: center; font-size: 20px">Output</p>"""
            )
            output_img = gr.Plot(label="Output Image",)
    with gr.Row():
        eval_btn = gr.Button(value="Evaluate")
    with gr.Row():
        gr.Markdown("## Examples")

    support_image.change(process_img,
                         inputs=[support_image, global_state],
                         outputs=[kp_support_image, global_state])
    kp_support_image.select(get_select_coords,
                       [global_state],
                       [global_state, kp_support_image],
                       queue=False,)
    confirm_kp_button.click(reset_skeleton,
                            inputs=global_state,
                            outputs=skel_support_image)
    undo_kp_button.click(reset_kp,
                         inputs=global_state,
                         outputs=[kp_support_image, skel_support_image])
    undo_skel_button.click(reset_skeleton,
                            inputs=global_state,
                            outputs=skel_support_image)
    skel_support_image.select(select_skeleton,
                              inputs=[global_state],
                              outputs=[global_state, skel_support_image])
    eval_btn.click(fn=process,
                   inputs=[query_image, global_state],
                   outputs=[output_img, global_state])

if __name__ == "__main__":
    print("Start app", parser.parse_args())
    gr.close_all()
    demo.launch(show_api=False)