import gradio as gr

from src.arguments import init_args
from src.GradioApp import GradioApp

app = GradioApp()

demo = gr.Blocks(**init_args.main_block)

with demo:
    with gr.Tabs() as tabs:

        # =============== #
        # Upload Data Tab #
        # =============== #
        with gr.TabItem('Upload Data', id=0):

            # Data uploader
            with gr.Row():

                with gr.Column():

                    data_requirements = gr.Markdown(
                        **init_args.data_requirements)

                    with gr.Row():
                        data_example_btn = gr.Button(
                            **init_args.data_example_btn)
                        data_example_b_btn = gr.Button(
                            **init_args.data_example_b_btn)

                data_uploader = gr.File(
                    **init_args.data_uploader)

            # Data Preview
            data_table = gr.Dataframe(
                **init_args.data_table)

            # Data visualization
            with gr.Column(**init_args.data_viz_col) as data_viz_col:

                data_viz_header = gr.Markdown(
                    '## Visualization')
                sku_plot_selector = gr.Dropdown(
                    **init_args.sku_plot_selector)
                data_viz_plot = gr.Plot(
                    **init_args.data_viz_plot)

                demand_fcst_btn = gr.Button(
                    **init_args.demand_fcst_btn)

        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
        # Upload Data Tab Event Listeners #
        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
        data_example_btn.click(
            app.data_example_btn__click,
            [],
            [data_table])

        data_example_b_btn.click(
            app.data_example_b_btn__click,
            [],
            [data_table])

        data_uploader.upload(
            app.data_uploader__upload,
            [data_uploader],
            [data_table])

        sku_plot_selector.change(
            app.sku_plot_selector__change,
            [sku_plot_selector],
            [data_viz_plot])

        # with gr.TabItem('Data Analytics', id=1):
        #     pass

        # ====================== #
        # Demand Forecasting Tab #
        # ====================== #
        with gr.TabItem('Demand Forecasting', id=2):

            gr.Markdown('# Select SKU for forecasting')

            with gr.Row():

                with gr.Column():

                    skus_to_fcst_dropdown = gr.Dropdown(
                        **init_args.skus_to_fcst_dropdown)
                    freq_options_radio = gr.Radio(
                        **init_args.freq_options_radio)
                    freq_convention_opt_radio = gr.Radio(
                        **init_args.freq_convention_opt_radio)
                    fcst_button = gr.Button(
                        **init_args.fcst_button)
                    freq_convention_opt_radio.change(
                        app.set_freq_convention,
                        [freq_convention_opt_radio],
                        [])

                with gr.Column():

                    fcst_data_prev_df = gr.Dataframe(
                        **init_args.fcst_data_prev_df)

            gr.Markdown('# View Forecasting Results')

            with gr.Row():
                # View Forecasted SKUS Dropdown
                skus_fcsted_dropdown = gr.Dropdown(
                    **init_args.skus_fcsted_dropdown)
                skus_fcsted_df = gr.Dataframe(
                    **init_args.skus_fcsted_df)

            skus_fcsted_plot = gr.Plot()

            inventory_toggle_btn = gr.Button(
                **init_args.inventory_toggle_btn)

        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
        # Demand Forecasting Tab Event Listeners #
        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
        fcst_button.click(
            app.fcst_button__click,
            [],
            [skus_fcsted_dropdown, skus_fcsted_df, skus_fcsted_plot])

        skus_fcsted_dropdown.change(
            app.skus_fcsted_dropdown__change,
            [skus_fcsted_dropdown],
            [skus_fcsted_df, skus_fcsted_plot])

        freq_options_radio.change(
            app.freq_options_radio__change, [freq_options_radio], [])

        # ========================== #
        # Inventory Optimization Tab #
        # ========================== #
        with gr.TabItem('Inventory Optimization', id=3) as io_tab:
            with gr.Row():
                gr.Markdown(**init_args.inventory_markdown)

                with gr.Column():
                    inventory_storage_capacity = gr.Number(
                        **init_args.inventory_storage_capacity)
                    inventory_budget_constraint = gr.Number(
                        **init_args.inventory_budget_constraint)

            inventory_dataframe = gr.Dataframe(
                **init_args.inventory_dataframe)

            with gr.Row():
                inventory_demo_button = gr.Button(
                    **init_args.inventory_demo_button)
                inventory_cta_button = gr.Button(
                    **init_args.inventory_cta_button)

            inventory_recommendations_md = gr.Markdown(
                **init_args.inventory_recommendations_md)

            inventory_recommendations_df = gr.Dataframe(
                **init_args.inventory_recommendations_df)
            inventory_viz = gr.Plot(
                **init_args.inventory_viz)

        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
        # Inventory Optimization Event Listeners #
        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
        inventory_storage_capacity.change(
            app.inventory_storage_capacity__change,
            [inventory_storage_capacity], [])

        inventory_budget_constraint.change(
            app.inventory_budget_constraint__change,
            [inventory_budget_constraint], []
        )

        inventory_cta_button.click(
            app.inventory_cta_button__click,
            [inventory_dataframe],
            [inventory_recommendations_md, inventory_recommendations_df, inventory_viz]
        )

        inventory_dataframe.input(
            app.inventory_dataframe__input,
            [inventory_dataframe], []
        )

        inventory_demo_button.click(
            app.inventory_demo_button__click,
            [],
            [inventory_dataframe, inventory_storage_capacity,
             inventory_budget_constraint]
        )

    # ~~~~~~~~~~~~~~~~~~~~~~ #
    # ~~~~~~~~~~~~~~~~~~~~~~ #
    # Global Event Listeners #
    # ~~~~~~~~~~~~~~~~~~~~~~ #
    # ~~~~~~~~~~~~~~~~~~~~~~ #
    demand_fcst_btn.click(
        app.demand_fcst_btn__click,
        None,
        tabs)

    inventory_toggle_btn.click(
        app.inventory_toggle_btn__click,
        None,
        tabs)

    data_table.change(
        app.data_table__change,
        [],
        [sku_plot_selector,
         skus_to_fcst_dropdown,
         data_viz_col,
         data_viz_plot])

    skus_to_fcst_dropdown.change(
        app.skus_to_fcst_dropdown__change,
        [skus_to_fcst_dropdown],
        [fcst_data_prev_df])

    io_tab.select(
        app.io_tab__select,
        [],
        [inventory_dataframe])
    
    gr.Markdown('Version 0.0.2 | Updated: 2024-Jan-31')

demo.launch()