import time
import shutil
import gradio as gr
from pygm_rrwm import pygm_rrwm


PYGM_IMG_DEFAULT_PATH = "media/pygm_default.png"
PYGM_SOLUTION_1_PATH = "media/pygm_image_1.png"
PYGM_SOLUTION_2_PATH = "media/pygm_image_2.png"


def _handle_pygm_solve(
    img_1_path: str,
    img_2_path: str,
    kpts1_path: str,
    kpts2_path: str,
):
    if img_1_path is None:
        raise gr.Error("Please upload file completely!")
    if img_2_path is None:
        raise gr.Error("Please upload file completely!")
    if kpts1_path is None:
        raise gr.Error("Please upload file completely!")
    if kpts1_path is None:
        raise gr.Error("Please upload file completely!")
    
    start_time = time.time()
    pygm_rrwm(
        img1_path=img_1_path,
        img2_path=img_2_path,
        kpts1_path=kpts1_path,
        kpts2_path=kpts2_path,
        output_path="media",
        filename="pygm_image"
    )
    solved_time = time.time() - start_time
    
    message = "Successfully solve the PYGM problem, using time ({:.3f}s).".format(solved_time)
    
    return message, PYGM_SOLUTION_1_PATH, PYGM_SOLUTION_2_PATH
    

def handle_pygm_solve(
    img_1_path: str,
    img_2_path: str,
    kpts1_path: str,
    kpts2_path: str,
):
    try:
        message = _handle_pygm_solve(
            img_1_path=img_1_path,
            img_2_path=img_2_path,
            kpts1_path=kpts1_path,
            kpts2_path=kpts2_path,
        )
        return message
    except Exception as e:
        message = str(e)
        return message, PYGM_SOLUTION_1_PATH, PYGM_SOLUTION_2_PATH


def handle_pygm_clear():
    shutil.copy(
        src=PYGM_IMG_DEFAULT_PATH,
        dst=PYGM_SOLUTION_1_PATH
    )
    shutil.copy(
        src=PYGM_IMG_DEFAULT_PATH,
        dst=PYGM_SOLUTION_2_PATH
    )

    message = "successfully clear the files!"
    return message, PYGM_SOLUTION_1_PATH, PYGM_SOLUTION_2_PATH


def convert_image_path_to_bytes(image_path):
    with open(image_path, "rb") as f:
        image_bytes = f.read()
    return image_bytes


with gr.Blocks() as pygm_page:

    gr.Markdown(
        '''
        This space displays the solution to the Graph Matching problem.
        ## How to use this Space?
        - Upload two pairs of '.png' and '.mat' files.
        - The results of the Graph Matching problem will be shown after you click the solve button.
        - Click the 'clear' button to clear all the files.
        ## Examples
        - You can get the test examples from our [Graph Match Dataset Repo.](https://huggingface.co/datasets/SJTU-TES/Graph-Match)
        '''
    )

    with gr.Row(variant="panel"):
        with gr.Column(scale=2):
            with gr.Row():
                pygm_img_1 = gr.File(
                    label="Upload .png File 1",
                    file_types=[".png"],
                    min_width=40,
                )
                pygm_img_2 = gr.File(
                    label="Upload .png File 2",
                    file_types=[".png"],
                    min_width=40,
                )
            with gr.Row():
                pygm_kpts_1 = gr.File(
                    label="Upload .mat File 1",
                    file_types=[".mat"],
                    min_width=40,
                )
                pygm_kpts_2 = gr.File(
                    label="Upload .mat File 2",
                    file_types=[".mat"],
                    min_width=40,
                )
            info = gr.Textbox(
                value="",
                label="Log",
                scale=4,
            )
        with gr.Column(scale=2):
            pygm_solution_1 = gr.Image(
                value=PYGM_SOLUTION_1_PATH, 
                type="filepath",
                label="Original Images"
            )
            pygm_solution_2 = gr.Image(
                value=PYGM_SOLUTION_2_PATH, 
                type="filepath", 
                label="Graph Matching Results"
            )
    with gr.Row():
        with gr.Column(scale=1, min_width=100):
            solve_button = gr.Button(
                value="Solve", 
                variant="primary", 
                scale=1
            )
        with gr.Column(scale=1, min_width=100):
            clear_button = gr.Button(
                "Clear", 
                variant="secondary", 
                scale=1
            )
        with gr.Column(scale=8):
            pass         
    
    solve_button.click(
        handle_pygm_solve,
        [pygm_img_1, pygm_img_2, pygm_kpts_1, pygm_kpts_2],
        outputs=[info, pygm_solution_1, pygm_solution_2]
    )
    
    clear_button.click(
        handle_pygm_clear,
        inputs=None,
        outputs=[info, pygm_solution_1, pygm_solution_2]
    )


if __name__ == "__main__":
    pygm_page.launch(debug = True)