import gradio as gr
import pandas as pd
import numpy as np
import joblib

# Load the model and scaler
model = joblib.load('catboost_model.pkl')
scaler = joblib.load('scaler.pkl')

# Define feature lists (make sure the order is consistent with training)
cat_features = ['Gender', 'Race', 'Education', 'Smoke', 'Alcohol']
num_features = ['Age', 'BMI', 'PIR', 'MPAH', 'PFDE', 'PFHxS', 'PFNA', 'PFOA', 'PFOS', 'PFUA']

# Define the prediction function
def predict_ra(gender, race, education, smoke, alcohol, age, bmi, pir, mpah, pfde, pfhxs, pfna, pfoa, pfos, pfua):
    """
    Predicts Rheumatoid arthritis risk based on input features.
    """
    # 1. Create a DataFrame
    input_data = pd.DataFrame({
        'Gender': [gender],
        'Race': [race],
        'Education': [education],
        'Smoke': [smoke],
        'Alcohol': [alcohol],
        'Age': [age],
        'BMI': [bmi],
        'PIR': [pir],
        'MPAH': [mpah],
        'PFDE': [pfde],
        'PFHxS': [pfhxs],
        'PFNA': [pfna],
        'PFOA': [pfoa],
        'PFOS': [pfos],
        'PFUA': [pfua]
    })

    # 2. Standardize numerical features
    input_data[num_features] = scaler.transform(input_data[num_features])

    # 3. Make a prediction
    prediction = model.predict(input_data)[0]
    probability = model.predict_proba(input_data)[0][1]

    # 4. Format the output
    if prediction == 1:
        result_text = "Prediction: Rheumatoid arthritis"
    else:
        result_text = "Prediction: Healthy"



    return f"{result_text}\nRisk Probability: {probability:.2%}"


# Define the Gradio interface using gr.Blocks() for layout flexibility
with gr.Blocks() as iface:
    gr.Markdown("# Rheumatoid arthritis Risk Prediction Calculator")  # Title
    gr.Markdown("Enter patient information to predict Rheumatoid arthritis risk.")  # Description

    with gr.Row():  # Layout: Use Rows and Columns for organization
        with gr.Column():
            gr.Markdown("## Basic Information")
            gender = gr.Radio([1, 2], label="Gender", info="1: Male, 2: Female", value=1)
            race = gr.Dropdown([1, 2, 3, 4, 5], label="Race", info="1: Mexican American, 2: Other Hispanic, 3: Non-Hispanic White, 4: Non-Hispanic Black, 5: Other Race", value=1)
            education = gr.Radio([0, 1], label="Education Level", info="0: High school or below, 1: Above high school", value=0)
            smoke = gr.Checkbox(label="Smoking", info="Does the patient smoke?")
            alcohol = gr.Checkbox(label="Alcohol Consumption", info="Does the patient consume alcohol?")

        with gr.Column():
            gr.Markdown("## Physical Indicators")
            age = gr.Number(label="Age", value=30)
            bmi = gr.Number(label="BMI", value=22.0)
            pir = gr.Number(label="PIR", value=2.0)
            mpah = gr.Number(label="MPAH (ng/mL)", value=0.5)
            pfde = gr.Number(label="PFDE (ng/mL)", value=0.5)
            pfhxs = gr.Number(label="PFHxS (ng/mL)", value=1.0)
            pfna = gr.Number(label="PFNA (ng/mL)", value=0.5)
            pfoa = gr.Number(label="PFOA (ng/mL)", value=2.0)
            pfos = gr.Number(label="PFOS (ng/mL)", value=5.0)
            pfua = gr.Number(label="PFUA (ng/mL)", value=0.5)

    predict_button = gr.Button("Predict")  # Prediction button
    output_textbox = gr.Textbox(label="Prediction Result") # Output text box

    predict_button.click(predict_ra,
                        inputs=[gender, race, education, smoke, alcohol, age, bmi, pir, mpah, pfde, pfhxs, pfna, pfoa, pfos, pfua],
                        outputs=output_textbox)

# Launch the Gradio app
iface.launch()