import gradio as gr
import pandas as pd

available_datasets = [
    "Demographic",
    "Health",
    "Financial",
    "Social",
    "Economic",
    "Political",
]

default_dataset = "Demographic"

available_attributes = [
    "Income",
    "Age",
    "Marital Status",
    "Education",
]

default_attributes = ["Income", "Age"]


def load_dataset(dataset):
    """
    Load and return data based on the 'dataset' parameter.

    Parameters:
    - dataset (str): The name of the dataset to load. Currently, this parameter is not used as the function returns a static DataFrame.

    Returns:
    - pd.DataFrame: A pandas DataFrame containing the loaded data.
    """
    return pd.DataFrame(
        {
            "Name": ["John", "Doe", "Jane", "Smith"],
            "Age": [25, 30, 35, 40],
            "Income": [50000, 60000, 70000, 80000],
            "Marital Status": ["Single", "Married", "Single", "Married"],
            "Education": ["High School", "Bachelor", "Master", "PhD"],
        }
    )


def predict(dataset, attributes, access_token):
    """
    Simulates predictions based on the provided dataset and attributes. Requires an access token for authentication.

    Parameters:
    - dataset (str): The name of the dataset on which to base predictions. This parameter is currently not used as the function returns a static prediction.
    - attributes (list of str): The attributes selected for prediction. This parameter is currently not used as the function returns a static prediction.
    - access_token (str): The access token required for authentication to make predictions.

    Returns:
    - tuple: A tuple containing a prediction message (str) and a pandas DataFrame with prediction results.

    Raises:
    - gr.Error: If the access token is not provided, an error is raised indicating that an access token is required.
    """
    if not access_token:
        raise gr.Error(
            "Access token missing or invalid. Please ensure you have a valid access token."
        )

    prediction_message = "200 predictions made in 2.5 seconds."
    prediction_results = pd.DataFrame(
        {
            "Name": ["John", "Doe", "Jane", "Smith"],
            "Predicted Value": [25000, 30000, 35000, 40000],
        }
    )
    return prediction_message, prediction_results


def load_dataset_and_predict(dataset, attributes, access_token):
    """
    Combines data loading and prediction into a single step, requiring an access token for the prediction part. Uses the 'load_dataset' and 'predict' functions to load the data and generate predictions.

    Parameters:
    - dataset (str): The name of the dataset to load and predict on.
    - attributes (list of str): The attributes selected for making predictions.
    - access_token (str): The access token required for authentication to make predictions.

    Returns:
    - tuple: A tuple containing the loaded pandas DataFrame, a prediction message (str), and a pandas DataFrame with prediction results.

    Note:
    - If the access token is not provided, the prediction part is skipped, and the prediction message and results will be returned as empty or None.
    """
    loaded_data = load_dataset(dataset)

    if access_token:
        prediction_message, prediction_results = predict(
            dataset, attributes, access_token
        )
    else:
        prediction_message = "No access token provided, prediction skipped."
        prediction_results = None

    return loaded_data, prediction_message, prediction_results


interface_theme = gr.themes.Default()

with gr.Blocks(theme=interface_theme) as demo:
    gr.Markdown("### Authenticate")
    access_token = gr.Textbox(
        type="password",
        label="Access Token",
        placeholder="Enter your access token here.",
    )

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Select Dataset and Attributes")

            selected_dataset = gr.Dropdown(
                choices=available_datasets,
                label="Select Dataset",
                value=default_dataset,
            )

            selected_attributes = gr.Dropdown(
                choices=available_attributes,
                label="Select Attributes",
                info="You can select multiple attributes.",
                multiselect=True,
                value=default_attributes,
            )

    gr.Markdown("### Dataset Preview")
    dataset_preview = gr.Dataframe()

    predict_button = gr.Button("Predict Attributes")

    gr.Markdown("### Prediction Results")
    prediction_preview = gr.Dataframe()
    prediction_label = gr.Markdown("")

    selected_dataset.change(
        fn=load_dataset, inputs=selected_dataset, outputs=dataset_preview
    )

    predict_button.click(
        fn=predict,
        inputs=[selected_dataset, selected_attributes, access_token],
        outputs=[prediction_label, prediction_preview],
    )

    demo.load(
        fn=load_dataset_and_predict,
        inputs=[selected_dataset, selected_attributes, access_token],
        outputs=[dataset_preview, prediction_label, prediction_preview],
    )

demo.launch()