from utils import (
    update_leaderboard_multilingual, 
    update_leaderboard_one_vs_all,
    handle_evaluation,
    process_results_file,
    create_html_image,
)
import os
import gradio as gr

from constants import *

if __name__ == "__main__":
    
    with gr.Blocks() as app:
               
        base_path = os.path.dirname(__file__)
        local_image_path = os.path.join(base_path, 'open_arabic_lid_arena.png')
        
        gr.HTML(create_html_image(local_image_path))
        gr.Markdown("# 🏅 Open Arabic Dialect Identification Leaderboard")
        
        # Multi-dialects leaderboard
        with gr.Tab("Multi-dialects model leaderboard"):
            gr.Markdown("""
                        Complete leaderboard across multiple arabic dialects. 
                        Compare the performance of different models across various metrics such as FNR, FPR, and other clasical metrics.
                        """
            )
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Select country to display")
                    country_selector = gr.Dropdown(
                        choices=supported_dialects,
                        value='Morocco', # Default to Morocco of course
                        label="Country"
                    )
                
                with gr.Column(scale=2):
                    gr.Markdown("### Select metrics to display")
                    metric_checkboxes = gr.CheckboxGroup(
                        choices=metrics,
                        value=default_metrics,
                        label="Metrics"
                    )
            
            with gr.Row():
                leaderboard_table = gr.DataFrame(
                    interactive=False
                )

            gr.Markdown("</br>")
            
            gr.Markdown("## Contribute to the Leaderboard")
            gr.Markdown("""
                        We welcome contributions from the community! 
                        If you have a model that you would like to see on the leaderboard, please use the 'Evaluate a model' or 'Upload your results' tabs to submit your model's performance.
                        Let's work together to improve Arabic dialect identification! 🚀
                        """)
            
        # Dialect confusion leaderboard
        with gr.Tab("Dialect confusion leaderboard"): # use to be "One-vs-All leaderboard"
            
            gr.Markdown("""
                        Detailed analysis of how well models distinguish specific dialects from others.
                        For each target dialect, see how often models incorrectly classify other dialects as the target.
                        Lower `false_positive_rate` indicate better ability to identify the true dialect by
                        showing **how often it misclassifies other dialects as the target dialect**.
                        """
            )
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Select your target language")
                    target_language_selector = gr.Dropdown(
                        choices=languages_to_display_one_vs_all,
                        value='Morocco', # Default to Morocco of course
                        label="Target Language"
                    )
                
                with gr.Column(scale=2):
                    gr.Markdown("### Select languages to compare to")
                    languages_checkboxes = gr.CheckboxGroup(
                        choices=languages_to_display_one_vs_all,
                        value=default_languages,
                        label="Languages"
                    )
                
            with gr.Row():
                binary_leaderboard_table = gr.DataFrame(
                    interactive=False
                )
            

        with gr.Tab("Evaluate a model"):
            gr.Markdown("Suggest a model to evaluate 🤗 (Supports only **Fasttext** models as SfayaLID, GlotLID, OpenLID, etc.)")
            gr.Markdown("For other models, you are welcome to **submit your results** through the upload section.")
            
            model_path = gr.Textbox(label="Model Path", placeholder='path/to/model')
            model_path_bin = gr.Textbox(label=".bin filename", placeholder='model.bin')
            gr.Markdown("### **⚠️ To ensure correct results, tick this when the model's labels are the iso_codes**")
            use_mapping = gr.Checkbox(label="Does not map to country", value=True)      # Initially enabled
            eval_button = gr.Button("Evaluate", value=False)                            # Initially disabled
            
            # Status message area
            status_message = gr.Markdown(value="")
            def update_status_message():
                return "### **⚠️Evaluating... Please wait...**"
            
            eval_button.click(update_status_message, outputs=[status_message])
            
            eval_button.click(handle_evaluation, inputs=[model_path, model_path_bin, use_mapping], outputs=[leaderboard_table, status_message])

        with gr.Tab("Upload your results"):

            # Define a code block to display
            code_snippet = """ 
            ```python
            
            # Load your model
            model = ... # Load your model here
            
            # Load evaluation benchmark
            eval_dataset = load_dataset("atlasia/Arabic-LID-Leaderboard", split='test').to_pandas() # do not change this line :)
            
            # Predict labels using your model
            eval_dataset['preds'] = eval_dataset['text'].apply(lambda text: predict_label(text, model)) # predict_label is a function that you need to define for your model
            
            # now drop the columns that are not needed, i.e. 'text', 'metadata' and 'dataset_source'
            df_eval = df_eval.drop(columns=['text', 'metadata', 'dataset_source'])
            df_eval.to_csv('your_model_name.csv')
            
            # submit your results: 'your_model_name.csv' to the leaderboard
            
            ```
            """
            
            gr.Markdown("## Upload your results to the leaderboard 🚀")
            gr.Markdown("### Submission guidelines: Run the test dataset on your model and save the results in a CSV file. Bellow a code snippet to help you with that.")
            gr.Markdown("### Nota Bene: The One-vs-All leaderboard evaluation is currently unavailable with the csv upload but will be implemented soon. Stay tuned!")
            gr.Markdown(code_snippet)

            uploaded_model_name = gr.Textbox(label="Model name", placeholder='Your model/team name')
            file = gr.File(label="Upload your results")
            upload_button = gr.Button("Upload")
            
            # Status message area
            status_message = gr.Markdown(value="")
            def update_status_message():
                return "### **⚠️Evaluating... Please wait...**"
            
            upload_button.click(update_status_message, outputs=[status_message])
            
            upload_button.click(process_results_file, inputs=[file, uploaded_model_name], outputs=[leaderboard_table, status_message])
        
        # Update multilangual table when any input changes
        country_selector.change(
            update_leaderboard_multilingual,
            inputs=[country_selector, metric_checkboxes],
            outputs=leaderboard_table
        )
        
        metric_checkboxes.change(
            update_leaderboard_multilingual,
            inputs=[country_selector, metric_checkboxes],
            outputs=leaderboard_table
        )
        
        # Update binary table when any input changes
        target_language_selector.change(
            update_leaderboard_one_vs_all,
            inputs=[target_language_selector, languages_checkboxes],
            outputs=[binary_leaderboard_table, languages_checkboxes]
        )
        
        languages_checkboxes.change(
            update_leaderboard_one_vs_all,
            inputs=[target_language_selector, languages_checkboxes],
            outputs=[binary_leaderboard_table, languages_checkboxes]
        )
        
        # Define load event to run at startup
        app.load(
            update_leaderboard_one_vs_all,
            inputs=[target_language_selector, languages_checkboxes],
            outputs=[binary_leaderboard_table, languages_checkboxes]
        )
        
        app.load(
            update_leaderboard_multilingual,
            inputs=[country_selector, metric_checkboxes],
            outputs=leaderboard_table
        )
        
    app.launch(allowed_paths=[base_path])