import gradio as gr
import pickle
import numpy as np

model = pickle.load(open('random_forest_model.pkl', 'rb'))

SEVERITY_LABELS = {0: 'Aucune', 1: 'Légère', 2: 'Modérée', 3: 'Sévère'}

def predict(age, gender, severity, test_last, test_mean, test_trend, previous_medication, prescribed_medication):
    features = np.array([[age, gender, severity, test_last, test_mean,
                          test_trend, previous_medication, prescribed_medication]])
    prediction = model.predict(features)[0]
    probabilities = model.predict_proba(features)[0]
    return SEVERITY_LABELS[prediction], {SEVERITY_LABELS[i]: float(p) for i, p in enumerate(probabilities)}

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Number(label='Age'),
        gr.Radio([0, 1], label='Gender (0=F, 1=M)'),
        gr.Slider(0, 3, step=1, label='Current Severity'),
        gr.Number(label='Last Test Score'),
        gr.Number(label='Mean Test Score'),
        gr.Number(label='Test Trend'),
        gr.Slider(0, 2, step=1, label='Previous Medication'),
        gr.Slider(0, 2, step=1, label='Prescribed Medication'),
    ],
    outputs=[
        gr.Text(label='Predicted Future Severity'),
        gr.Label(label='Class Probabilities'),
    ],
    title='Alzheimer Severity Predictor',
)

demo.launch()