import gradio as gr
from df.PaperCentral import PaperCentral
from gradio_calendar import Calendar
from datetime import datetime, timedelta
from typing import Union, List

# Initialize the PaperCentral class instance
paper_central_df = PaperCentral()

# Create the Gradio Blocks app with custom CSS
with gr.Blocks(css="style.css") as demo:
    gr.Markdown("# Paper Central")

    # Create a row for navigation buttons and calendar
    with gr.Row():
        with gr.Column(scale=1):
            # Define the 'Next Day' and 'Previous Day' buttons
            next_day_btn = gr.Button("Next Day")
            prev_day_btn = gr.Button("Previous Day")
        with gr.Column(scale=4):
            # Define the calendar component for date selection
            calendar = Calendar(
                type="datetime",
                label="Select a date",
                info="Click the calendar icon to bring up the calendar.",
                value=datetime.today().strftime('%Y-%m-%d')  # Default to today's date
            )

    # Create a row for Hugging Face options and Conference options
    with gr.Row():
        with gr.Column():
            # Define the checkbox group for Hugging Face options
            cat_options = gr.CheckboxGroup(
                label="Category",
                choices=[
                    'cs.*',
                    'eess.*',
                    'econ.*',
                    'math.*',
                    'astro-ph.*',
                    'cond-mat.*',
                    'gr-qc',
                    'hep-ex',
                    'hep-lat',
                    'hep-ph',
                    'hep-th',
                    'math-ph',
                    'nlin.*',
                    'nucl-ex',
                    'nucl-th',
                    'physics.*',
                    'quant-ph',
                    'q-bio.*',
                    'q-fin.*',
                    'stat.*',
                ],
                value=["cs.*"]
            )
            hf_options = gr.CheckboxGroup(
                label="Hugging Face options",
                choices=["🤗 paper-page", "datasets", "models", "spaces", "github"],
                elem_id="hf_options"
            )

        with gr.Column():
            # Define the checkbox group for Conference options
            conference_options = gr.CheckboxGroup(
                label="Conference options",
                choices=["In proceedings"] + PaperCentral.CONFERENCES
            )

            # Define a Textbox for author search
            author_search = gr.Textbox(
                label="Search Authors",
                placeholder="Enter author name",
            )
    # Define the Dataframe component to display paper data
    # List of columns in your DataFrame
    columns = paper_central_df.COLUMNS_START_PAPER_PAGE

    paper_central_component = gr.Dataframe(
        label="Paper Data",
        value=paper_central_df.df_prettified[columns],
        datatype=[
            paper_central_df.DATATYPES[column]
            for column in columns
        ],
        row_count=(0, "dynamic"),
        interactive=False,
        height=1000,
        elem_id="table",
    )


    # Define function to move to the next day
    def go_to_next_day(
            date: Union[str, datetime],
            cat_options_list: List[str],
            hf_options_list: List[str],
            conference_options_list: List[str],
            author_search_input: str,
    ) -> tuple:
        """
        Moves the selected date to the next day and updates the data.

        Args:
            date (Union[str, datetime]): The current date selected in the calendar.
            cat_options_list (List[str]): List of selected Category options.
            hf_options_list (List[str]): List of selected Hugging Face options.
            conference_options_list (List[str]): List of selected Conference options.

        Returns:
            tuple: The new date as a string and the updated Dataframe component.
        """
        # Ensure the date is in string format
        if isinstance(date, datetime):
            date_str = date.strftime('%Y-%m-%d')
        else:
            date_str = date

        # Parse the date string and add one day
        new_date = datetime.strptime(date_str, '%Y-%m-%d') + timedelta(days=1)
        new_date_str = new_date.strftime('%Y-%m-%d')

        # Update the Dataframe
        updated_data = paper_central_df.filter(
            selected_date=new_date_str,
            cat_options=cat_options_list,
            hf_options=hf_options_list,
            conference_options=conference_options_list,
            author_search_input=author_search_input,
        )

        # Return the new date and updated Dataframe
        return new_date_str, updated_data


    # Define function to move to the previous day
    def go_to_previous_day(
            date: Union[str, datetime],
            cat_options_list: List[str],
            hf_options_list: List[str],
            conference_options_list: List[str],
            author_search_input: str,
    ) -> tuple:
        """
        Moves the selected date to the previous day and updates the data.

        Args:
            date (Union[str, datetime]): The current date selected in the calendar.
            cat_options_list (List[str]): List of selected Category options.
            hf_options_list (List[str]): List of selected Hugging Face options.
            conference_options_list (List[str]): List of selected Conference options.

        Returns:
            tuple: The new date as a string and the updated Dataframe component.
        """
        # Ensure the date is in string format
        if isinstance(date, datetime):
            date_str = date.strftime('%Y-%m-%d')
        else:
            date_str = date

        # Parse the date string and subtract one day
        new_date = datetime.strptime(date_str, '%Y-%m-%d') - timedelta(days=1)
        new_date_str = new_date.strftime('%Y-%m-%d')

        # Update the Dataframe
        updated_data = paper_central_df.filter(
            selected_date=new_date_str,
            cat_options=cat_options_list,
            hf_options=hf_options_list,
            conference_options=conference_options_list,
            author_search_input=author_search_input,
        )

        # Return the new date and updated Dataframe
        return new_date_str, updated_data


    # Define function to update data when date or options change
    def update_data(
            date: Union[str, datetime],
            cat_options_list: List[str],
            hf_options_list: List[str],
            conference_options_list: List[str],
            author_search_input: str,
    ):
        """
        Updates the data displayed in the Dataframe based on the selected date and options.

        Args:
            date (Union[str, datetime]): The selected date.
            cat_options_list (List[str]): List of selected Category options.
            hf_options_list (List[str]): List of selected Hugging Face options.
            conference_options_list (List[str]): List of selected Conference options.

        Returns:
            gr.Dataframe.update: An update object for the Dataframe component.
        """
        return paper_central_df.filter(
            selected_date=date,
            cat_options=cat_options_list,
            hf_options=hf_options_list,
            conference_options=conference_options_list,
            author_search_input=author_search_input,
        )


    # Function to handle conference options change
    def on_conference_options_change(
            date: Union[str, datetime],
            cat_options_list: List[str],
            hf_options_list: List[str],
            conference_options_list: List[str],
            author_search_input: str,
    ):

        cat_options_update = gr.update()
        paper_central_component_update = gr.update()
        visible = True

        # Some conference options are selected
        # Update cat_options to empty list
        if conference_options_list:
            cat_options_update = gr.update(value=[])
            paper_central_component_update = update_data(
                date,
                [],
                hf_options_list,
                conference_options_list,
                author_search_input,
            )
            visible = False

        calendar_update = gr.update(visible=visible)
        next_day_btn_update = gr.update(visible=visible)
        prev_day_btn_update = gr.update(visible=visible)

        return paper_central_component_update, cat_options_update, calendar_update, next_day_btn_update, prev_day_btn_update


    # Function to handle category options change
    def on_cat_options_change(
            date: Union[str, datetime],
            cat_options_list: List[str],
            hf_options_list: List[str],
            conference_options_list: List[str],
            author_search_input: str,
    ):
        conference_options_update = gr.update()
        paper_central_component_update = gr.update()
        visible = False

        # Some category options are selected
        # Update conference_options to empty list
        if cat_options_list:
            conference_options_update = gr.update(value=[])
            paper_central_component_update = update_data(
                date,
                cat_options_list,
                hf_options_list,
                [],
                author_search_input,
            )
            visible = True

        calendar_update = gr.update(visible=visible)
        next_day_btn_update = gr.update(visible=visible)
        prev_day_btn_update = gr.update(visible=visible)

        return paper_central_component_update, conference_options_update, calendar_update, next_day_btn_update, prev_day_btn_update


    inputs = [
        calendar,
        cat_options,
        hf_options,
        conference_options,
        author_search,
    ]

    # Set up the event listener for the author search
    author_search.submit(
        fn=update_data,
        inputs=inputs,
        outputs=paper_central_component,
    )


    # Set up the event listener for the 'Next Day' button
    next_day_btn.click(
        fn=go_to_next_day,
        inputs=inputs,
        outputs=[calendar, paper_central_component],
    )

    # Set up the event listener for the 'Previous Day' button
    prev_day_btn.click(
        fn=go_to_previous_day,
        inputs=inputs,
        outputs=[calendar, paper_central_component],
    )

    # Set up the event listener for the calendar date change
    calendar.change(
        fn=update_data,
        inputs=inputs,
        outputs=paper_central_component,
    )

    # Set up the event listener for the Hugging Face options change
    hf_options.change(
        fn=update_data,
        inputs=inputs,
        outputs=paper_central_component,
    )


    # Event chaining for conference options change
    conference_options.change(
        fn=on_conference_options_change,
        inputs=inputs,
        outputs=[paper_central_component, cat_options, calendar, next_day_btn, prev_day_btn],
    )

    # Event chaining for category options change
    cat_options.change(
        fn=on_cat_options_change,
        inputs=inputs,
        outputs=[paper_central_component, conference_options, calendar, next_day_btn, prev_day_btn],
    )

    # Load the initial data when the app starts
    demo.load(
        fn=update_data,
        inputs=inputs,
        outputs=paper_central_component,
        api_name=False,
    )


# Define the main function to launch the app
def main():
    """
    Launches the Gradio app.
    """
    demo.launch()


# Run the main function when the script is executed
if __name__ == "__main__":
    main()