import json
import os
import os.path as osp
import gradio as gr
import numpy as np
import gradio as gr


def load_json(load_dir_path, json_file_name):
    
    load_path = os.path.join(load_dir_path, json_file_name)
    if not os.path.exists(load_path):
        return None
    with open(load_path, 'r', encoding='utf-8') as f:
        obj_serializable = json.load(f)
    return obj_serializable


def load_results_recaption(save_path, model="gpt-3.5-turbo-0125"):
    result_list = load_json(save_path, f'final_results-{model}.json')
    if result_list is not None:
        result_list = result_list['result_list']
    
    if result_list is None:
        result_list = load_json(save_path, 'inference_results.json')   

    return result_list

plava_theme = gr.themes.Monochrome(
    text_size="sm",
    spacing_size="sm",
    primary_hue=gr.themes.Color(c100="#f5f5f5", c200="#e5e5e5", c300="#d4d4d4", c400="#a3a3a3", c50="#fafafa", c500="#737373", c600="#525252", c700="#404040", c800="#262626", c900="#171717", c950="#000000"),
    secondary_hue=gr.themes.Color(c100="#f5f5f5", c200="#e5e5e5", c300="#d4d4d4", c400="#a3a3a3", c50="#fafafa", c500="#737373", c600="#525252", c700="#404040", c800="#262626", c900="#171717", c950="#000000"),
    neutral_hue=gr.themes.Color(c100="#f5f5f5", c200="#e5e5e5", c300="#d4d4d4", c400="#a3a3a3", c50="#fafafa", c500="#737373", c600="#525252", c700="#404040", c800="#262626", c900="#171717", c950="#000000"),
).set(
    background_fill_primary_dark='*primary_950',
    background_fill_secondary_dark='*neutral_950'
)




load_results_funcs = [
    load_results_recaption,
]


recaption_root_dir = "recaption_results"
local_video_root_dir = "DATAS/Recaption/Inter4K/60fps/UHD"
remote_video_root_dir ="https://huggingface.co/datasets/ermu2001/Inter4kPllavaRecaption/resolve/main/DATAS/Recaption/Inter4K/60fps/UHD"

def show(result_list_first, result_list_second, result_index):
    sample2index_second = {}

    for i, result in enumerate(result_list_second):
        if 'video_path' not in result:
            continue

        question = result['question'] if 'question' in result else ''
        video_path = result['video_path']
        samplehash = question + '--' +video_path
        sample2index_second[samplehash] = i

    info = result_list_first[result_index]
    info_str_first = json.dumps(info, indent=4, ensure_ascii=False)
    video_path = info['video_path']
    question = info['question'] if 'question' in info else ''
    samplehash = question + '--' +video_path
    if samplehash in sample2index_second:
        info = result_list_second[sample2index_second[samplehash]]
        info_str_second = json.dumps(info, indent=4, ensure_ascii=False)
    else:
        info_str_second = f"NO {video_path} IN THE SECOND RESULT DIR"
    video_path = video_path.replace(local_video_root_dir, remote_video_root_dir)
    return video_path, info_str_first, info_str_second

def reload_results_dirs():
    result_dirs = []
    # load result dir paths
    for dirpath, dirnames, filenames in os.walk(recaption_root_dir):
        if len(dirnames) == 0 and len(filenames) != 0:
            result_dirs.append(dirpath)
    return gr.Dropdown(result_dirs, value=result_dirs[0])

def reload_results(result_dir):
    # if isinstance(result_dir, list):
    #     result_dir = result_dir[0]

    if result_dir is None or not osp.exists(result_dir):
        return None
    
    for fn in load_results_funcs:
        result_list = fn(result_dir)
        if result_list is not None:
            np.random.shuffle(result_list)
            break
    result_index = gr.Slider(0, len(result_list), step=1)

    return result_list, result_index



with gr.Blocks(title="PLAVA RESULTS", theme=plava_theme) as demo:
    result_list_first = gr.State()
    result_list_second = gr.State()

    with gr.Row():
        with gr.Column():
            gr.Markdown("# Showing off Model's Outputs.")
            gr.Markdown(
                "You can find all our results, including:\n"
                "1. results of Captioned Inter4k\n"
                "2. results of Different Benchmark inference outputs.\n"
                "Choose a directory to see the different output variant.\n"
                "You can also choose secondary directory (as long as they are from the same dataset.) to compare on the results.\n"
            )

    with gr.Row():
        with gr.Column():
            show_video = gr.Video(interactive=False)

        with gr.Column():
            button_reload = gr.Button(value='Reload From The Evaluation/Inference Root Directory')
            result_index = gr.Slider(0, 0, step=1, label="Index")

            result_dir_first = gr.Dropdown(label='Test Result Path')
            info_first = gr.Text(interactive=False, label='Detailed Output Information')
            result_dir_second = gr.Dropdown(label='Test Result Path')
            info_second = gr.Text(interactive=False, label='Detailed Output Information')
        

    button_reload.click(reload_results_dirs, [], [result_dir_first])
    button_reload.click(reload_results_dirs, [], [result_dir_second])
    result_dir_first.change(reload_results, [result_dir_first], [result_list_first, result_index])
    result_dir_second.change(reload_results, [result_dir_second], [result_list_second, result_index])
    result_index.change(show, [result_list_first, result_list_second, result_index], [show_video, info_first, info_second])
    demo.load(reload_results_dirs, [], [result_dir_first])
    demo.load(reload_results_dirs, [], [result_dir_second])
    
demo.launch()