import gradio as gr
import joblib
import pandas as pd
import datasets
import json

# Load the model
pipe = joblib.load("./model.pkl")

title = "Premium Amount Prediction"
description = "This model predicts the Premium Amount. Drag and drop any slice from the dataset or edit values as you wish in the dataframe component below."

# Load and prepare dataset
df = datasets.load_dataset("silvaKenpachi/mental_health")["train"].to_pandas()
df.dropna(axis=0, inplace=True)

# Load configuration
with open("./config.json") as f:
    config_dict = json.load(f)
all_headers = config_dict["sklearn"]["columns"]

# Filter headers to only include those present in the dataset
headers = [col for col in all_headers if col in df.columns]

# Define input and output interfaces
inputs = [gr.Dataframe(headers=headers, row_count=(2, "dynamic"), col_count=(len(headers), "fixed"), label="Input Data", interactive=True)]
outputs = [gr.Dataframe(row_count=(2, "dynamic"), col_count=(1, "fixed"), label="Predictions", headers=["Depression"])]

def infer(inputs):
    data = pd.DataFrame(inputs, columns=headers)
    predictions = pipe.predict(data)
    return pd.DataFrame(predictions, columns=["Depression"])

gr.Interface(
    fn=infer,
    inputs=inputs,
    outputs=outputs,
    title=title,
    description=description,
    examples=[df[headers].head(3).values.tolist()],
    cache_examples=False
).launch(debug=True)