import gradio as gr
import xgboost as xgb
import pandas as pd
from datasets import load_dataset

# Load the dataset
dataset = load_dataset("Ammok/hair_health")

# Convert to Pandas DataFrame for exploration
df = pd.DataFrame(dataset['train'])

# Example: Train a simple XGBoost model
X = df.drop(columns=["target_column"])  # Replace with your feature columns
y = df["target_column"]  # Replace with your target column

# Train a basic XGBoost model (replace with custom model training code)
model = xgb.XGBClassifier()
model.fit(X, y)

# Function for making predictions
def predict(input_data):
    data = pd.DataFrame([input_data], columns=X.columns)
    prediction = model.predict(data)
    return prediction[0]

# Set up Gradio interface for data exploration
def explore_data(row_number):
    return df.iloc[row_number].to_dict()

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# Hair Health Dataset Exploration")
    
    row_number_input = gr.Number(label="Row Number")
    data_output = gr.JSON(label="Row Data")
    row_number_input.change(explore_data, inputs=[row_number_input], outputs=[data_output])

    gr.Markdown("## Make a Prediction")
    input_data = {col: gr.Number(label=col) for col in X.columns}  # Adjust based on features
    output = gr.Textbox(label="Prediction")
    
    submit_button = gr.Button("Predict")
    submit_button.click(predict, inputs=[input_data], outputs=[output])

demo.launch()