import gradio as gr
import joblib
import numpy as np
import pandas as pd
from sklearn.datasets import load_breast_cancer

# Load the pipeline
pipeline = joblib.load('knn_pipeline.joblib')

# Load the data to get feature names
data = load_breast_cancer()
feature_names = data.feature_names


def predict_cancer(*input_data):
    # Convert input data to a DataFrame with feature names
    input_array = np.array([input_data])
    input_df = pd.DataFrame(input_array, columns=feature_names)
    # Use the pipeline to predict
    prediction = pipeline.predict(input_df)
    # Map prediction to target names
    target_names = ['malignant', 'benign']
    return target_names[prediction[0]]

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            col1 = [gr.Slider(label=feature, minimum=0, maximum=10000, step=0.1) for feature in feature_names[:10]]
        with gr.Column():
            col2 = [gr.Slider(label=feature, minimum=0, maximum=10000, step=0.1) for feature in feature_names[10:20]]
        with gr.Column():
            col3 = [gr.Slider(label=feature, minimum=0, maximum=10000, step=0.1) for feature in feature_names[20:]]

    outputs = gr.Textbox(label="Prediction")
    btn = gr.Button("Predict")
    btn.click(fn=predict_cancer, inputs=col1 + col2 + col3, outputs=outputs)

demo.launch()