# Necessary imports
import gradio as gr
import pandas as pd
from pycaret.classification import load_model, predict_model

# Load the tuned model
tuned_gbc_classifier = load_model('tuned_gbc_classifier')

def predict_ten_year_chd(male, age, education, currentSmoker, cigsPerDay, BPMeds, prevalentStroke,
                         prevalentHyp, diabetes, totChol, sysBP, diaBP, BMI, heartRate, glucose):
    try:
        # Convert categorical variables to numerical representation
        male = 1 if male == "Male" else 0
        education_mapping = {
            "Some High School": 0,
            "High School Graduate": 1,
            "Some College": 2,
            "College Graduate": 3
        }
        education = education_mapping.get(education, 0)

        # Create a DataFrame with the input values
        data = pd.DataFrame(
            data=[[male, age, education, currentSmoker, cigsPerDay, BPMeds, prevalentStroke,
                   prevalentHyp, diabetes, totChol, sysBP, diaBP, BMI, heartRate, glucose]],
            columns=['male', 'age', 'education', 'currentSmoker', 'cigsPerDay', 'BPMeds', 'prevalentStroke',
                     'prevalentHyp', 'diabetes', 'totChol', 'sysBP', 'diaBP', 'BMI', 'heartRate', 'glucose']
        )

        # Make a prediction
        pred = predict_model(tuned_gbc_classifier, data=data)

        # Extract the prediction and the confidence using the correct keys
        prediction = pred['prediction_label'].iloc[0]
        confidence = pred['prediction_score'].iloc[0]

        # Return the prediction with 'At Risk' category for No CHD with confidence < 0.8
        if prediction == 0 and confidence < 0.8:
            return f"Prediction: No CHD (At Risk), Confidence: {confidence:.2f}"
        else:
            return f"Prediction: {'Has CHD' if prediction == 1 else 'No CHD'}, Confidence: {confidence:.2f}"

    except Exception as e:
        return f"An error occurred: {str(e)}"

# Create the Gradio interface
iface = gr.Interface(
    fn=predict_ten_year_chd,
    inputs=[
        gr.inputs.Radio(["Male", "Female"], label="Gender"),
        gr.inputs.Slider(minimum=18, maximum=100, label="Age"),
        gr.inputs.Dropdown(["Some High School", "High School Graduate", "Some College", "College Graduate"], label="Education"),
        gr.inputs.Checkbox(label="Current Smoker"),
        gr.inputs.Slider(minimum=0, maximum=50, default=0, label="Cigarettes Per Day"),
        gr.inputs.Checkbox(label="On Blood Pressure Medication"),
        gr.inputs.Checkbox(label="History of Prevalent Stroke"),
        gr.inputs.Checkbox(label="History of Prevalent Hypertension"),
        gr.inputs.Checkbox(label="Diabetes"),
        gr.inputs.Slider(minimum=100, maximum=400, default=200, label="Total Cholesterol"),
        gr.inputs.Slider(minimum=90, maximum=200, default=120, label="Systolic BP"),
        gr.inputs.Slider(minimum=60, maximum=120, default=80, label="Diastolic BP"),
        gr.inputs.Slider(minimum=15, maximum=50, default=25, label="BMI"),
        gr.inputs.Slider(minimum=40, maximum=120, default=75, label="Heart Rate"),
        gr.inputs.Slider(minimum=40, maximum=300, default=100, label="Glucose Level")
    ],
    outputs=gr.outputs.Textbox(),
    live=False,  # set live to False to add a submit button
    title="CHD Prediction",
    description="By Abderrahim Benmoussa, Ph.D."
)

# Run the app
iface.launch()