ui/layout.py

import gradio as gr
import time
from logic.supabase_client import auth_handler

from logic.data_utils import CustomHFDatasetSaver
from data.lang2eng_map import lang2eng_mapping
from gradio_modal import Modal
from logic.handlers import *
from logic.image_transforms import rotate_image_90_left, rotate_image_90_right, reset_image_to_original
from config.settings import *
from functools import partial

from .selection_page import build_selection_page
from .main_page import build_main_page
from .main_page import sort_with_pyuca

js_code = """
function() {
    // Get the full URL with the fragment
    const url = window.location.href;
    const fragment = url.split('#')[1];

    if (!fragment) {
        return "";
    }

    // Parse the fragment into an object
    const params = new URLSearchParams(fragment);
    const access_token = params.get('access_token');
    const refresh_token = params.get('refresh_token');

    // Create a JSON string with the tokens
    const tokens = JSON.stringify({
        access_token: access_token,
        refresh_token: refresh_token
    });

    // Return the JSON string to the Gradio output component
    return tokens;
}
"""

scroll_to_top_js = """
function() {
    // Scroll to the top of the page smoothly
    window.scrollTo({
        top: 0,
        behavior: 'smooth'
    });
    return "";
}
"""

def login_user(email, password):
    result = auth_handler.login(email, password)
    if result['success']:
        session_data = result['data']
        persistent_data = {
            "refresh_token": session_data['refresh_token'],
            "user_email": session_data['user_email']
        }
        return session_data['client'], persistent_data, result['message']
    else:
        persistent_data = {
            "refresh_token": "",
            "user_email": ""
        }
        return None, persistent_data, result['message']

def login_user_recovery(session_data: str):
    """
    This function receives session data (tokens as a JSON string) from the frontend,
    retrieves the session, and returns data in a format similar to login_user.
    """
    try:
        import json
        tokens = json.loads(session_data)
        access_token = tokens.get("access_token")
        refresh_token = tokens.get("refresh_token")

        if not access_token or not refresh_token:
            return None, gr.skip(), "Invalid session data provided."

        result = auth_handler.retrieve_session_from_tokens(access_token, refresh_token)
        
        if result['success']:
            session_data_result = result['data']
            persistent_data = {
                "refresh_token": session_data_result['refresh_token'],
                "user_email": session_data_result['user_email']
            }
            return session_data_result['client'], persistent_data, result['message']
        else:
            persistent_data = {
                "refresh_token": "",
                "user_email": ""
            }
            return None, persistent_data, result['message']

    except Exception as e:
        return None, gr.skip(), f"Failed to process recovery login: {e}"

def sign_up(email, password):
    result = auth_handler.sign_up(email, password)
    return result['message']

def reset_password(email):
    result = auth_handler.reset_password_for_email(email)
    return result['message']

def log_out(supabase_user_client, persistent_session):
    """
    Logs out the user and clears the session. If error occurs, it returns an empty persistent session (logging out user).
    """
    persistent_session = {
        "refresh_token": "",
        "user_email": ""
    }
    if supabase_user_client:
        result = auth_handler.logout(supabase_user_client)
        if result['success']:
            print("User logged out successfully.")
            return persistent_session
        else:
            print(f"Error logging out: {result['message']}")
            return persistent_session
    else:
        print("No user client provided to log out.")
        return persistent_session

def restore_user_session(session_data, login_status=None):
    print("Restoring user session with data:", session_data)
    # defualt values if the user is not logged in
    # or the session data is not valid
    login_status_update = gr.update(value= login_status if login_status else "")
    proceed_button_update = gr.update(value="Proceed as Anonymous User", interactive=True)
    login_button_update = gr.update(visible=True)
    sign_up_button_update = gr.update(visible=True)
    reset_password_button_update = gr.update(visible=True)
    logout_button_update = gr.update(visible=False)
    change_password_field_update = gr.update(visible=False)
    change_password_field_confirm_update = gr.update(visible=False)
    change_password_button_update = gr.update(visible=False)
    change_password_status_update = gr.update(value="")
    persistent_data = {
        "refresh_token": "",
        "user_email": ""
    }
    if not session_data or not session_data.get('refresh_token', ''):
        print("No session data found, proceeding as anonymous user.")
        return None, persistent_data, login_status_update, proceed_button_update, login_button_update, sign_up_button_update, reset_password_button_update, logout_button_update, change_password_field_update, change_password_field_confirm_update, change_password_button_update, change_password_status_update
    
    result = auth_handler.restore_session(session_data['refresh_token'])
    if result['success']:
        restored_session = result['data']
        new_persistent_data = {
            "refresh_token": restored_session['refresh_token'],
            "user_email": restored_session['user_email']
        }
        login_status_update = gr.update(value=result['message'])
        proceed_button_update = gr.update(value="Proceed", interactive=True)
        login_button_update = gr.update(visible=False)
        sign_up_button_update = gr.update(visible=False)
        reset_password_button_update = gr.update(visible=False)
        logout_button_update = gr.update(visible=True)
        change_password_field_update = gr.update(visible=True)
        change_password_field_confirm_update = gr.update(visible=True)
        change_password_button_update = gr.update(visible=True)
        return restored_session['client'], new_persistent_data, login_status_update, proceed_button_update, login_button_update, sign_up_button_update, reset_password_button_update, logout_button_update, change_password_field_update, change_password_field_confirm_update, change_password_button_update, change_password_status_update
    else:
        return None, persistent_data, login_status_update, proceed_button_update, login_button_update, sign_up_button_update, reset_password_button_update, logout_button_update, change_password_field_update, change_password_field_confirm_update, change_password_button_update, change_password_status_update

def change_password(supabase_user_client, new_password, confirm_password):
    """
    Changes the user's password.
    """
    if new_password != confirm_password:
        return "Passwords do not match. Please try again."
    result = auth_handler.change_password(supabase_user_client, new_password)
    return result['message']


def get_key_by_value(dictionary, value):
    for key, val in dictionary.items():
        if val == value:
            return key
    return None

def load_concepts(category, concept_btn, local_storage, loading_example, concepts):
    country, lang, _, _ = local_storage
    if category:
        if lang in words_mapping:
            category = get_key_by_value(words_mapping[lang], category) # translate back the category name
        concept_val = None
        # I'm doing this to avoid the case where the user load a previous example and the concept 'reamins' selected even if category changed
        if loading_example:
            concept_val = concept_btn if concept_btn else None
        if country is None or lang is None:
            return gr.update(choices=[], value=None), False
        else:
            return gr.update(choices=sort_with_pyuca(concepts[country][lang2eng_mapping.get(lang, lang)][category]), value=concept_val), False
    else:
        concepts_list = []
        # cats = []  # FIXME: Assumes concepts are unique across all categories
        if country is None or lang is None:
            return gr.update(choices=[], value=None), False
        else:
            for cat in concepts[country][lang]:
                for concept in concepts[country][lang][cat]:
                    # cats.append(cat)
                    concepts_list.append(concept)
            return gr.update(choices=sort_with_pyuca(concepts_list), value=None), False
        # return gr.update(choices=[], value=None), False

def update_category(category, concept_btn, local_storage, loading_example, concepts):
    country, lang, _, _ = local_storage
    if not category and concept_btn:
        # FIXME: Assumes unique concepts across categories
        for cat in concepts[country][lang2eng_mapping.get(lang, lang)]:
            for concept in concepts[country][lang2eng_mapping.get(lang, lang)][cat]:
                if concept == concept_btn:
                    category = cat
        return gr.update(choices=sort_with_pyuca(concepts[country][lang2eng_mapping.get(lang, lang)]), value=category), True
    else:
        # import pdb; pdb.set_trace()
        if country is None or lang is None:
            return gr.update(choices=[], value=None), False
        else:
            return gr.update(choices=sort_with_pyuca(concepts[country][lang2eng_mapping.get(lang, lang)]), value=category), False


def switch_ui(country, language, username, password, flag=False, metadata_dict=None):
    # wait for 2 seconds
    print(f"Language: {language}, Country: {country}")

    if not country or not language:
        gr.Warning(f"⚠️ Please select Country and Language first.",  duration=3)
        return

    # Hide the selection page and show the main UI.
    if flag:
        gr.Info(f"Loading the Language Selection UI",  duration=3)
        time.sleep(2)
        local_storage = [None, None, "", ""]
        countries = sort_with_pyuca(list(metadata_dict.keys()))
        return (local_storage, gr.update(visible=flag), gr.update(visible=not flag), 
                gr.update(visible=not flag), 
                gr.update(choices=countries, label="Country", value=None), 
                gr.update(value=None, choices=[], label="Language", allow_custom_value=False, interactive=False))
    
    else:
        gr.Info(f"Loading the Main UI",  duration=3)
        time.sleep(2)
        local_storage = [country, language, username, password]
        return local_storage, gr.update(visible=flag), gr.update(visible=not flag), gr.update(visible=not flag)

def build_ui(concepts_dict, metadata_dict, HF_API_TOKEN, HF_DATASET_NAME):
    hf_writer = CustomHFDatasetSaver(HF_API_TOKEN, HF_DATASET_NAME, private=True)
    custom_css = """
        .compact-container {
            max-width: 600px;
            margin: auto;
            padding: 20px;
        }
        .compact-btn {
        }
        #image_inp img {
            object-fit: contain;        /* make sure the full image shows */
            height: 460px;           /* set a fixed height */
        }
        #vlm_output .input-container {
        position: relative;
        }
        #vlm_output .input-container::before {
        content: "";
        position: absolute;
        top: 0; left: 0; right: 0; bottom: 0;
        z-index: 10;             /* sits above the textarea */
        background: transparent;
        }
    """
    ############################################################################
    with gr.Blocks(css=custom_css) as ui:
        supabase_user_client = gr.State(None)
        persistent_session = gr.BrowserState(None)

        local_storage = gr.State([None, None, "", ""])
        loading_example = gr.State(False) # to check if the values are loaded from a user click on an example in 
        # First page: selection

        selection_page, country_choice, language_choice, proceed_btn, username, password, intro_markdown, login_btn, sign_up_btn, reset_password_btn, login_status, logout_btn, change_password_field, change_password_field_confirm, change_password_btn, change_password_status = build_selection_page(metadata_dict)
        
        # Second page
        cmp_main_ui = build_main_page(concepts_dict, metadata_dict, local_storage)

        main_ui_placeholder = cmp_main_ui["main_ui_placeholder"]
        country_inp = cmp_main_ui["country_inp"]
        language_inp = cmp_main_ui["language_inp"]
        image_inp = cmp_main_ui["image_inp"]
        image_url_inp = cmp_main_ui["image_url_inp"]
        long_caption_inp = cmp_main_ui["long_caption_inp"]
        num_words_inp = cmp_main_ui["num_words_inp"]
        category_btn = cmp_main_ui["category_btn"]
        concept_btn = cmp_main_ui["concept_btn"]
        category_concept_dropdowns = cmp_main_ui["category_concept_dropdowns"]
        instruct_btn = cmp_main_ui["instruct_btn"]
        clear_btn = cmp_main_ui["clear_btn"]
        hide_faces_btn = cmp_main_ui["hide_faces_btn"]
        hide_all_faces_btn = cmp_main_ui["hide_all_faces_btn"]
        unhide_faces_btn = cmp_main_ui["unhide_faces_btn"]
        submit_btn = cmp_main_ui["submit_btn"]
        timestamp_btn = cmp_main_ui["timestamp_btn"]
        exampleid_btn = cmp_main_ui["exampleid_btn"]
        username_inp = cmp_main_ui["username_inp"]
        password_inp = cmp_main_ui["password_inp"]
        modal_saving = cmp_main_ui["modal_saving"]
        modal_data_saved = cmp_main_ui["modal_data_saved"]
        modal = cmp_main_ui["modal"]
        exit_btn = cmp_main_ui["exit_btn"]
        intro_text_inp = cmp_main_ui["intro_text_inp"] # Note intro_text is actually the "Task" text from metadata
        intro_text_inst_inp = cmp_main_ui["intro_text_inst_inp"] # This is the "Instructions" text from metadata
        modal_saving_text = cmp_main_ui["modal_saving_text"]
        modal_data_saved_text = cmp_main_ui["modal_data_saved_text"]
        exclude_btn = cmp_main_ui["exclude_btn"]
        modal_exclude_confirm = cmp_main_ui["modal_exclude_confirm"]
        cancel_exclude_btn = cmp_main_ui["cancel_exclude_btn"]
        confirm_exclude_btn = cmp_main_ui["confirm_exclude_btn"]
        vlm_output = cmp_main_ui["vlm_output"]
        gen_button = cmp_main_ui["gen_button"]
        vlm_feedback = cmp_main_ui["vlm_feedback"]
        modal_vlm = cmp_main_ui["modal_vlm"]
        vlm_no_btn = cmp_main_ui["vlm_no_btn"]
        vlm_done_btn = cmp_main_ui["vlm_done_btn"]
        submit_yes = cmp_main_ui["submit_yes"]
        submit_no = cmp_main_ui["submit_no"]
        modal_submit = cmp_main_ui["modal_submit"]
        vlm_cancel_btn = cmp_main_ui["vlm_cancel_btn"]
        vlm_model_dropdown = cmp_main_ui["vlm_model_dropdown"]
        
        # Image rotation buttons
        rotate_left_btn = cmp_main_ui["rotate_left_btn"]
        rotate_right_btn = cmp_main_ui["rotate_right_btn"]
        reset_image_btn = cmp_main_ui["reset_image_btn"]

        # dictionary to store ALL vlm_outputs and vlm_models by exampleid
        vlm_captions = gr.State(None)
        vlm_models = gr.State(None)
        ### Category button
        category_btn.change(
            fn=partial(load_concepts, concepts=concepts_dict),
            inputs=[category_btn, concept_btn, local_storage, loading_example],
            outputs=[concept_btn, loading_example]
        )
        concept_btn.change(
            fn=partial(update_category, concepts=concepts_dict),
            inputs=[category_btn, concept_btn, local_storage, loading_example],
            outputs=[category_btn, loading_example]
        ).then(
            fn=partial(load_concepts, concepts=concepts_dict),
            inputs=[category_btn, concept_btn, local_storage, gr.State(True)],
            outputs=[concept_btn, loading_example]
        )
        ####

        # is_blurred = gr.State(False)  # Initialize as False

        # Notification image is false
        with Modal(visible=False) as modal_img_url:
            gr.Markdown("The image URL is not valid. Please provide a valid image URL.")
            # gr.Warning(f"⚠️ The image URL is not valid. Please provide a valid image URL.")

        # Event listeners
        gr.on(
            triggers=[country_choice.change, language_choice.change],
            fn=validate_metadata,
            inputs=[country_choice, language_choice],
            outputs=[proceed_btn],
        )
        gr.on(
            triggers=[image_url_inp.change],
            fn=update_image,
            inputs=[image_url_inp],
            outputs=[image_inp, modal_img_url]
        )
        gr.on(
            triggers=[language_choice.change],
            fn=partial(update_intro_language, metadata=metadata_dict),
            inputs=[country_choice, language_choice, intro_markdown],
            outputs=[intro_markdown]
        )

        gr.on(
            triggers=[image_inp.change, long_caption_inp.change],
            fn=update_timestamp,
            outputs=[timestamp_btn]
        )
        ori_img = gr.State(None)
        gr.on(
            triggers=[image_inp.change],
            fn=validate_inputs,
            inputs=[image_inp, ori_img, concept_btn], # is_blurred
            outputs=[submit_btn, image_inp, ori_img], # is_blurred
        )

        gr.on(
            triggers=[long_caption_inp.change],
            fn=count_words,
            inputs=[long_caption_inp, language_choice],
            outputs=[num_words_inp],
        )

        # Clear Button
        clear_btn.click(
            fn=clear_data,
            outputs=[
                image_inp, image_url_inp, long_caption_inp, vlm_output, vlm_feedback, vlm_done_btn, vlm_no_btn, gen_button, exampleid_btn, 
                category_btn, concept_btn, 
                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                category_concept_dropdowns[3], category_concept_dropdowns[4]
            ],
        )

        #============= Face Blurring ============= #
        with Modal(visible=False) as modal_faces:
            with gr.Column():
                face_img = gr.Image(label="Image Faces", elem_id="image_faces", format="png", height=512, width=768)
                with gr.Row():
                    faces_count = gr.Textbox(label="Face Counts", elem_id="face_counts", interactive=False)
                    blur_faces_ids = gr.Dropdown(
                        [], value=[], multiselect=True, label="Please select the faces IDs you want to blur.", elem_id="blur_faces_ids")
                    # blur_faces_ids = gr.Textbox(label="Specify faces ids to blur by comma", elem_id="blur_faces_ids", interactive=True)
                    submit_btn_face = gr.Button("Submit", variant="primary", interactive=True, elem_id="submit_btn_face")
        
        with Modal(visible=False) as modal_faces_2:
            gr.Markdown("The current image does not have any faces to hide.")
    
        faces_info = gr.State(None)
        hide_faces_btn.click(
            fn=select_faces_to_hide,
            inputs=[image_inp, blur_faces_ids],
            outputs=[image_inp, modal_faces, modal_faces_2, face_img, faces_count, faces_info, blur_faces_ids]
        )

        submit_btn_face.click(
            fn=blur_selected_faces,
            inputs=[image_inp, blur_faces_ids, faces_info, face_img, faces_count], # is_blurred 
            outputs=[image_inp, modal_faces, face_img, faces_count, blur_faces_ids] # is_blurred
        )

        hide_all_faces_btn.click(
            fn=blur_all_faces,
            inputs=[image_inp],
            outputs=[image_inp, modal_faces_2] # is_blurred
        )

        unhide_faces_btn.click(
            fn=unhide_faces,
            inputs=[image_inp, ori_img], # is_blurred
            outputs=[image_inp] # is_blurred
        )

        # ============= Image Rotation ============= #
        rotate_left_btn.click(
            fn=rotate_image_90_left,
            inputs=[image_inp],
            outputs=[image_inp]
        )

        rotate_right_btn.click(
            fn=rotate_image_90_right,
            inputs=[image_inp],
            outputs=[image_inp]
        )

        reset_image_btn.click(
            fn=reset_image_to_original,
            inputs=[image_inp, ori_img],
            outputs=[image_inp]
        )

        # ===============================
        with gr.Column(visible=False, elem_id="browse_data") as browse_data_placeholder:
            # Browse Data
            browse_data_text = gr.Markdown("## Browse Data")

            loading_msg = gr.Markdown("**Loading your data, please wait ...**")
            for contrib_type in ["Your data", "Other data"]:
                with gr.Tab(contrib_type):
                    if contrib_type == "Your data":
                        # Show user's past data points
                        user_examples = gr.Dataset(
                            samples=[],
                            components=['image','textbox','textbox','textbox','textbox', 
                                       'textbox','textbox','textbox', 'textbox'],
                            headers=['Image', 'Image URL (Optional, if not uploading an image)', 'Description', 'Country', 'Language', 
                                    'Category', 'Concept', 'Additional Concepts', 'ID'],
                        )
                        # Handle clicking on an example
                        user_examples.click(
                            fn=partial(handle_click_example, concepts_dict=concepts_dict),
                            inputs=[user_examples, vlm_captions, vlm_models],
                            outputs=[
                                image_inp, image_url_inp, long_caption_inp, exampleid_btn, 
                                category_btn, concept_btn,
                                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                                category_concept_dropdowns[3], category_concept_dropdowns[4], loading_example, vlm_output, vlm_model_dropdown
                            ],
                        )

                    elif contrib_type == "Other data":
                        # Show others' images
                        pass

        # ============================================ #
        # Submit Button Click events
        login_btn.click(
            fn=login_user,
            inputs=[username, password],
            outputs=[supabase_user_client, persistent_session, login_status],
        ).then(
            fn=restore_user_session,
            inputs=[persistent_session, login_status],
            outputs=[supabase_user_client, persistent_session, login_status, proceed_btn, login_btn, sign_up_btn, reset_password_btn, logout_btn, change_password_field, change_password_field_confirm, change_password_btn, change_password_status],
        )

        sign_up_btn.click(
            fn=sign_up,
            inputs=[username, password],
            outputs=[login_status],
        )

        logout_btn.click(
            fn=log_out,
            inputs=[supabase_user_client, persistent_session],
            outputs=[persistent_session]
        ).then(
            fn=restore_user_session,
            inputs=[persistent_session],
            outputs=[supabase_user_client, persistent_session, login_status, proceed_btn, login_btn, sign_up_btn, reset_password_btn, logout_btn, change_password_field, change_password_field_confirm, change_password_btn, change_password_status],
        )
        change_password_btn.click(
            fn=change_password,
            inputs=[supabase_user_client, change_password_field, change_password_field_confirm],
            outputs=[change_password_status]
        )
        reset_password_btn.click(
            fn=reset_password,
            inputs=[username],
            outputs=[login_status]
        )

        proceed_btn.click(
            fn=partial(switch_ui, flag=False),
            inputs=[country_choice, language_choice, username, password],
            outputs=[local_storage, selection_page, main_ui_placeholder, browse_data_placeholder],
        ).then(
            fn=partial(update_language, metadata_dict=metadata_dict, concepts_dict=concepts_dict),
            inputs=[local_storage],
            outputs=[
                country_inp, language_inp, username_inp, password_inp, category_btn, concept_btn, image_inp, 
                image_url_inp, long_caption_inp, intro_text_inp, intro_text_inst_inp, instruct_btn, clear_btn,
                submit_btn, modal_saving_text, modal_data_saved_text, timestamp_btn, exit_btn, browse_data_text,
                loading_msg, hide_all_faces_btn, hide_faces_btn, unhide_faces_btn, exclude_btn,
                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                category_concept_dropdowns[3], category_concept_dropdowns[4]
            ]
        ).then(
            fn=partial(update_user_data, HF_DATASET_NAME=HF_DATASET_NAME, local_ds_directory_path = LOCAL_DS_DIRECTORY_PATH),
            inputs=[supabase_user_client, country_choice, language_choice],
            outputs=[user_examples, loading_msg, vlm_captions, vlm_models],
        )


        # Exit Button
        exit_btn.click(
            fn=exit,
            outputs=[
                image_inp, image_url_inp, long_caption_inp, vlm_output, vlm_feedback, vlm_done_btn, vlm_no_btn, gen_button, user_examples, loading_msg, 
                username, password, local_storage, exampleid_btn, category_btn, concept_btn,
                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                category_concept_dropdowns[3], category_concept_dropdowns[4]
            ],
        ).success(
            fn=partial(switch_ui, flag=True, metadata_dict=metadata_dict),
            inputs=[country_choice, language_choice, username, password],
            outputs=[local_storage, selection_page, main_ui_placeholder, browse_data_placeholder, country_choice, language_choice],
        )

        # Disable button while saving
        # Note: I think this is not longer needed as we clear all the inputs and disable the submit button if the data is saved correctly
        # def disable_submit():
        #     return gr.update(interactive=False)

        # def enable_submit():
        #     return gr.update(interactive=True)

        # STEP 1: show modal
        # submit_btn.click(lambda: Modal(visible=True), None, modal_saving)
        
        # STEP 2: disable button
        # submit_btn.click(disable_submit, None, [submit_btn], queue=False)

        #STEP 3: perform save_data
        data_outputs = {
            "image": image_inp,
            "image_url": image_url_inp,
            "caption": long_caption_inp,
            "country": country_inp,
            "language": language_inp,
            "category": category_btn,
            "concept": concept_btn,
            "category_1_concepts": category_concept_dropdowns[0],
            "category_2_concepts": category_concept_dropdowns[1],
            "category_3_concepts": category_concept_dropdowns[2],
            "category_4_concepts": category_concept_dropdowns[3],
            "category_5_concepts": category_concept_dropdowns[4],
            "timestamp": timestamp_btn,
            "username": username_inp,
            "password": password_inp,
            "id": exampleid_btn,
            "excluded": gr.State(value=False),
            "concepts_dict": gr.State(value=concepts_dict),
            "country_lang_map": gr.State(value=lang2eng_mapping),
            "client": supabase_user_client,
            # "is_blurred": is_blurred
            "vlm_caption": vlm_output,
            "vlm_feedback": vlm_feedback,
            "vlm_model": vlm_model_dropdown # selected vlm model
        }
        # data_outputs = [image_inp, image_url_inp, long_caption_inp,
        #         country_inp, language_inp, category_btn, concept_btn,
        #         timestamp_btn, username_inp, password_inp, exampleid_btn]
        hf_writer.setup(list(data_outputs.keys()), local_ds_folder = LOCAL_DS_DIRECTORY_PATH)

        # STEP 4: Chain save_data, then update_user_data, then re-enable button, hide modal, and clear
        # submit_btn.click(lambda: Modal(visible=True), None, modal_vlm)
        submit_btn.click(submit_button_clicked,
                        inputs=[vlm_output],
                        outputs=[modal_vlm, modal_submit],
                        js=scroll_to_top_js)

        # submit_btn.click(partial(submit_button_clicked, save_fn=hf_writer.save, 
        #                             data_outputs=data_outputs),
        #                 inputs=[vlm_output], 
        #                 outputs=[modal_vlm, image_inp, image_url_inp, long_caption_inp, vlm_output, vlm_feedback, exampleid_btn, 
        #                          category_btn, concept_btn, category_concept_dropdowns[0], category_concept_dropdowns[1], 
        #                          category_concept_dropdowns[2], category_concept_dropdowns[3], category_concept_dropdowns[4]])
        
        def wire_submit_chain(button, modal_ui):
            e = button.click(
                fn=lambda: Modal(visible=False),
                outputs=[modal_ui]
            ).success(
                hf_writer.save,
                inputs = list(data_outputs.values()),
                outputs = None,
            ).success(
                fn=partial(clear_data, "submit"),
                outputs=[
                    image_inp, image_url_inp, long_caption_inp, vlm_output, vlm_feedback, vlm_done_btn, vlm_no_btn, gen_button, exampleid_btn, 
                    category_btn, concept_btn,
                    category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                    category_concept_dropdowns[3], category_concept_dropdowns[4]
                ],
            # ).success(enable_submit, 
            #    None, [submit_btn]
            # ).success(lambda: Modal(visible=False), 
            #    None, modal_saving
            # ).success(lambda: Modal(visible=True), 
            #        None, modal_data_saved
            ).success(
                # set loading msg
                lambda: gr.update(value="**Loading your data, please wait ...**"),
                None, loading_msg
            ).success(
                fn=partial(update_user_data, HF_DATASET_NAME=HF_DATASET_NAME, local_ds_directory_path = LOCAL_DS_DIRECTORY_PATH),
                inputs=[supabase_user_client, country_choice, language_choice],
                outputs=[user_examples, loading_msg, vlm_captions, vlm_models]
            )
            return e
        
        wire_submit_chain(vlm_done_btn, modal_vlm)
        wire_submit_chain(vlm_no_btn, modal_vlm)
        wire_submit_chain(submit_yes, modal_submit)
        submit_no.click(lambda: Modal(visible=False), None, modal_submit)
        vlm_cancel_btn.click(lambda: Modal(visible=False), None, modal_vlm)
        # ============================================ #
        # instructions button 
        instruct_btn.click(lambda: Modal(visible=True), None, modal)
        # ============================================ #

        # # Load saved values from local storage (browser storage)
        # @ui.load(inputs=[local_storage], outputs=[country_choice, language_choice, username, password])
        # def load_from_local_storage(saved_values):
        #     print("loading from local storage", saved_values)
        #     return saved_values[0], saved_values[1], saved_values[2], saved_values[3]
        
        # ============================================= #
        # Exclude button 
        # ============================================= #

        exclude_status = gr.Markdown(visible=False)

        # Show confirmation modal when exclude button is clicked
        exclude_btn.click(
            fn=check_exclude_fn,
            inputs=[image_inp],
            outputs=[modal_exclude_confirm]
        )

        # Close modal when cancel button is clicked
        cancel_exclude_btn.click(
            fn=lambda: gr.update(visible=False),
            outputs=[modal_exclude_confirm]
        )

        # Handle confirm exclusion
        # Create a modified data_outputs for exclusion with excluded=True
        exclude_data_outputs = data_outputs.copy()
        exclude_data_outputs["excluded"] = gr.State(value=True)
        
        confirm_exclude_btn.click(
            fn=lambda: gr.update(visible=False),
            outputs=[modal_exclude_confirm]
        ).success(
            fn=hf_writer.save,
            inputs=list(exclude_data_outputs.values()),
            outputs=None
        ).success(
            fn=partial(clear_data, "remove"),
            outputs=[ 
                image_inp, image_url_inp, long_caption_inp, vlm_output, vlm_feedback, vlm_done_btn, vlm_no_btn, gen_button, exampleid_btn, 
                category_btn, concept_btn, 
                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                category_concept_dropdowns[3], category_concept_dropdowns[4]
            ]
        ).success(
            fn=lambda: gr.update(value="Example excluded successfully", visible=True),
            outputs=exclude_status
        ).success(
            fn=lambda: gr.update(value="**Refreshing your data, please wait...**"),
            outputs=loading_msg
        ).success(
            fn=partial(update_user_data, HF_DATASET_NAME=HF_DATASET_NAME, local_ds_directory_path=LOCAL_DS_DIRECTORY_PATH),
            inputs=[supabase_user_client, country_choice, language_choice],
            outputs=[user_examples, loading_msg, vlm_captions, vlm_models]
        )
        # ============================================= #
        # VLM Gen button
        # ============================================= #
        gen_button.click(
            fn=generate_vlm_caption, # processor=processor, model=model
            inputs=[image_inp, vlm_model_dropdown],
            outputs=[vlm_output, vlm_feedback, vlm_done_btn, vlm_no_btn, gen_button, vlm_model_dropdown],
            concurrency_limit=1,
            concurrency_id="vlm_queue"
        )
        # vlm_output.change(
        #     fn=lambda : gr.update(interactive=False) if vlm_output.value else gr.update(interactive=True),
        #     inputs=[],
        #     outputs=[gen_button]
        # )

        ui.load(
            fn=login_user_recovery,
            inputs=gr.Textbox(visible=False, value=""), # hidden textbox to get the url tokens
            outputs=[supabase_user_client, persistent_session, login_status],
            js=js_code
        ).then(
            fn=restore_user_session,
            inputs=[persistent_session],
            outputs=[supabase_user_client, persistent_session, login_status, proceed_btn, login_btn, sign_up_btn, reset_password_btn, logout_btn, change_password_field, change_password_field_confirm, change_password_btn, change_password_status],
        )
    return ui