import yaml
import gradio as gr
import pandas as pd
from functools import partial
from nist_cda_dashboard.component import ComponentHelper
from nist_cda_dashboard.utils import DownloadHelper


with open("./config.yml", 'r') as file:
    config = yaml.safe_load(file)

component_helper = ComponentHelper(file_qc_results=pd.read_csv(config["file_path"]["file_qc_results"]),
                                   dataset_qc_results=pd.read_csv(
                                       config["file_path"]["dataset_qc_results"]),
                                   analyzed_gating_results=pd.read_csv(
                                       config["file_path"]["analyzed_gating_results"]),
                                   analyzed_gating_results_wCDAResults=pd.read_csv(
                                       config["file_path"]["analyzed_gating_results_wCDAResults"]),
                                   config=config)

custom_css = """
#plot_wScrollBar {
    overflow-x: auto !important;
    text-align: center !important;
}

#plot_wScrollBar > .js-plotly-plot,
#plot_wScrollBar > .plotly,
#plot_wScrollBar > div:first-child { 
    display: inline-block !important;
    text-align: left !important;
}
"""

with gr.Blocks(css=custom_css) as demo:
    with gr.Sidebar(position="left", width=700):
        gr.Markdown("# STEP 1: Filter Dataset")
        clear_all_models_button, select_all_models_button, model_filter, \
            clear_all_datasets_button, select_all_datasets_button, dataset_filter, \
            main_apply_filters_button = component_helper.creator.main_filter_components()

    clear_all_models_button.click(fn=partial(component_helper.updater.clean_all_choices),
                                  outputs=model_filter)
    select_all_models_button.click(fn=partial(component_helper.updater.select_all_choices, "model_filter"),
                                   outputs=model_filter)

    clear_all_datasets_button.click(fn=partial(component_helper.updater.clean_all_choices),
                                    outputs=dataset_filter)
    select_all_datasets_button.click(fn=partial(component_helper.updater.select_all_choices, "dataset_filter"),
                                     outputs=dataset_filter)

    model_filter.change(fn=component_helper.updater.update_dataset_filter,
                        inputs=[model_filter,
                                dataset_filter],
                        outputs=dataset_filter)

    gr.Markdown("# STEP 2: Choose QC or Analysis Tab")
    with gr.Tab(label="Quality Check"):
        gr.Markdown("# STEP 3: Filter File")
        clear_all_sop_exps_qc_tab_button, select_all_sop_exps_qc_tab_button, sop_exp_qc_tab_filter, \
            clear_all_materials_button, select_all_materials_button, material_filter, \
            clear_all_issues_button, select_all_issues_button, issue_filter, \
            qc_tab_apply_filters_button = component_helper.creator.qc_tab_filter_components()

        clear_all_sop_exps_qc_tab_button.click(fn=partial(component_helper.updater.clean_all_choices),
                                               outputs=sop_exp_qc_tab_filter)
        select_all_sop_exps_qc_tab_button.click(fn=partial(component_helper.updater.select_all_choices, "sop_exp_qc_tab_filter"),
                                                outputs=sop_exp_qc_tab_filter)

        clear_all_materials_button.click(fn=partial(component_helper.updater.clean_all_choices),
                                         outputs=material_filter)
        select_all_materials_button.click(fn=partial(component_helper.updater.select_all_choices, "material_filter"),
                                          outputs=material_filter)

        clear_all_issues_button.click(fn=partial(component_helper.updater.clean_all_choices),
                                      outputs=issue_filter)
        select_all_issues_button.click(fn=partial(component_helper.updater.select_all_choices, "issue_filter"),
                                       outputs=issue_filter)

        sop_exp_qc_tab_filter.change(fn=component_helper.updater.update_material_filter,
                                     inputs=[sop_exp_qc_tab_filter,
                                             material_filter],
                                     outputs=material_filter)

        gr.Markdown("<br><br>")
        gr.Markdown("# STEP 4: Results")
        qc_tab_dataset_qc_status_filter = component_helper.creator.qc_tab_dataset_qc_status_filter_component()

        with gr.Row():
            gr.Markdown("## File QC Results Table")
            download_file_qc_button = gr.Button(
                "Download File QC Results")

        file_qc_status_filter, file_qc_table, file_qc_table_no_file_msg = component_helper.creator.file_qc_result_components()

        download_file_qc_button_hidden = gr.DownloadButton(
            visible=False, elem_id="download_file_qc_button_hidden")
        download_file_qc_button.click(fn=(lambda df: DownloadHelper.df2csv(df, "file_qc_table")),
                                      inputs=[file_qc_table],
                                      outputs=[download_file_qc_button_hidden]). \
            then(fn=None,
                 inputs=None,
                 outputs=None,
                 js="() => document.querySelector('#download_file_qc_button_hidden').click()")

        gr.Markdown("<br><br>")
        with gr.Row():
            gr.Markdown("## Dataset QC Results Table")
            download_dataset_qc_button = gr.Button(
                "Download Dataset QC Results")
        dataset_qc_table, dataset_qc_table_no_dataset_msg = component_helper.creator.dataset_qc_result_components()

        download_dataset_qc_button_hidden = gr.DownloadButton(
            visible=False, elem_id="download_dataset_qc_button_hidden")
        download_dataset_qc_button.click(fn=(lambda df: DownloadHelper.df2csv(df, "dataset_qc_table")),
                                         inputs=[dataset_qc_table],
                                         outputs=[download_dataset_qc_button_hidden]). \
            then(fn=None,
                 inputs=None,
                 outputs=None,
                 js="() => document.querySelector('#download_dataset_qc_button_hidden').click()")

        gr.Markdown("<br><br>")
        with gr.Row():
            gr.Markdown("## QC Results Visualization")
            download_qc_fig_button = gr.Button(
                "Download QC Visualization Figure")
        qc_fig, qc_fig_no_file_msg = component_helper.creator.qc_result_visual_components()

        download_qc_fig_button_hidden = gr.DownloadButton(
            visible=False, elem_id="download_qc_fig_button_hidden")
        download_qc_fig_button.click(fn=(lambda fig: DownloadHelper.fig2png(fig, "qc_visual_fig")),
                                     inputs=[qc_fig],
                                     outputs=[download_qc_fig_button_hidden]). \
            then(fn=None,
                 inputs=None,
                 outputs=None,
                 js="() => document.querySelector('#download_qc_fig_button_hidden').click()")

        gr.on(fn=component_helper.updater.update_file_qc_table,
              triggers=[main_apply_filters_button.click, qc_tab_apply_filters_button.click,
                        qc_tab_dataset_qc_status_filter.change, file_qc_status_filter.change],
              inputs=[dataset_filter, sop_exp_qc_tab_filter, material_filter,
                      issue_filter, qc_tab_dataset_qc_status_filter, file_qc_status_filter],
              outputs=[file_qc_table, file_qc_table_no_file_msg,
                       dataset_qc_table, dataset_qc_table_no_dataset_msg])

        gr.on(fn=component_helper.updater.update_qc_fig,
              triggers=[main_apply_filters_button.click, qc_tab_apply_filters_button.click,
                        qc_tab_dataset_qc_status_filter.change],
              inputs=[dataset_filter, sop_exp_qc_tab_filter, material_filter,
                      issue_filter, qc_tab_dataset_qc_status_filter],
              outputs=[qc_fig, qc_fig_no_file_msg])

    with gr.Tab(label="Gating Result Analysis"):
        gr.Markdown("# STEP 3: Filter File")
        sample_filter, sop_exp_analysis_tab_filter, compensation_control_filter, \
            gating_control_filter, pop_pheno_parent_filter, clear_gating_tab_filters_button, analysis_tab_apply_filters_button = component_helper.creator.analysis_tab_filter_components()

        gr.Markdown("<br><br>")
        gr.Markdown("# STEP 4: Results")
        gr.Markdown("## Analysis Visualization")
        analysis_tab_dataset_qc_status_filter, analyzed_result_filter = component_helper.creator.analyzed_result_filter_component()

        with gr.Tab(label="Single Result  Barplot"):
            download_barplot_fig_button = gr.Button("Download Barplot Figure")
            analysis_barplot_fig, barplot_not_reportable_msg = component_helper.creator.analysis_barplot_components()

            download_barplot_fig_button_hidden = gr.DownloadButton(
                visible=False, elem_id="download_barplot_fig_button_hidden")
            download_barplot_fig_button.click(fn=(lambda fig: DownloadHelper.fig2png(fig, "analysis_barplot_fig")),
                                              inputs=[analysis_barplot_fig],
                                              outputs=[download_barplot_fig_button_hidden]). \
                then(fn=None,
                    inputs=None,
                    outputs=None,
                     js="() => document.querySelector('#download_barplot_fig_button_hidden').click()")

        with gr.Tab(label="Multiple Results Heatmap"):
            download_heatmap_fig_button = gr.Button(
                "Download Heatmap Figure")
            compared_protocol_filter, include_CDA_results_checkbox, analysis_heatmap_exp_info_table, analysis_heatmap_exp_comparison_table, \
                analysis_heatmap_fig, heatmap_not_reportable_msg = component_helper.creator.analysis_heatmap_components()

            download_heatmap_fig_button_hidden = gr.DownloadButton(
                visible=False, elem_id="download_heatmap_fig_button_hidden")
            download_heatmap_fig_button.click(fn=(lambda fig: DownloadHelper.fig2png(fig, "analysis_heatmap_fig")),
                                              inputs=[analysis_heatmap_fig],
                                              outputs=[download_heatmap_fig_button_hidden]). \
                then(fn=None,
                     inputs=None,
                     outputs=None,
                     js="() => document.querySelector('#download_heatmap_fig_button_hidden').click()")

        gr.on(fn=component_helper.updater.update_barplot_fig,
              triggers=[main_apply_filters_button.click, analysis_tab_apply_filters_button.click,
                        analysis_tab_dataset_qc_status_filter.change, analyzed_result_filter.change],
              inputs=[dataset_filter, analysis_tab_dataset_qc_status_filter,
                      sample_filter, sop_exp_analysis_tab_filter, compensation_control_filter,
                      gating_control_filter, pop_pheno_parent_filter, analyzed_result_filter],
              outputs=[analysis_barplot_fig, barplot_not_reportable_msg])

        gr.on(fn=component_helper.updater.update_heatmap_fig,
              triggers=[main_apply_filters_button.click, analysis_tab_apply_filters_button.click,
                        analysis_tab_dataset_qc_status_filter.change, analyzed_result_filter.change,
                        compared_protocol_filter.change, include_CDA_results_checkbox.change],
              inputs=[dataset_filter, analysis_tab_dataset_qc_status_filter,
                      sample_filter, sop_exp_analysis_tab_filter, compensation_control_filter,
                      gating_control_filter, pop_pheno_parent_filter,
                      analyzed_result_filter, compared_protocol_filter, include_CDA_results_checkbox],
              outputs=[analysis_heatmap_exp_info_table, analysis_heatmap_exp_comparison_table,
                       analysis_heatmap_fig, heatmap_not_reportable_msg])

        gr.on(fn=component_helper.updater.update_analysis_tab_filters,
              triggers=[sample_filter.change,
                        sop_exp_analysis_tab_filter.change, compensation_control_filter.change,
                        gating_control_filter.change, pop_pheno_parent_filter.change],
              inputs=[sample_filter,
                      sop_exp_analysis_tab_filter, compensation_control_filter,
                      gating_control_filter, pop_pheno_parent_filter],
              outputs=[sample_filter,
                       sop_exp_analysis_tab_filter, compensation_control_filter,
                       gating_control_filter, pop_pheno_parent_filter])

        clear_gating_tab_filters_button.click(fn=component_helper.updater.clear_analysis_tab_filters,
                                              outputs=[sample_filter,
                                                       sop_exp_analysis_tab_filter, compensation_control_filter,
                                                       gating_control_filter, pop_pheno_parent_filter])

demo.launch(inbrowser=True)