import gradio as gr
import pandas as pd
import joblib

# Load the preprocessor and model
preprocessor = joblib.load('preprocessor_v1.0.joblib')
model = joblib.load('calibrated_random_forest_model_updated_v1.1.joblib')

def predict(age, sex, race, etiology_cirrosis, hepatorenal_syndrome, omeprazole, spironolactone, furosemide, propanolol, dialisis,
            portal_vein_thrombosis, ascitis, hepatocellular_carcinoma, albumin, total_bilirrubin, direct_bilirrubina, inr, creatinine,
            platelets, ast, alt, hemoglobin, hematocrit, leucocytes, sodium, potassium, varices, red_wale_marks, rupture_point, 
            active_bleeding, therapy, terlipressin_dose, time_to_endoscophy_hours, rebleeding):
    input_data = {
        'age': age,
        'sex': sex,
        'race': race,
        'etiology_cirrosis': etiology_cirrosis,
        'hepatorenal_syndrome': hepatorenal_syndrome,
        'omeprazole': omeprazole,
        'spironolactone': spironolactone,
        'furosemide': furosemide,
        'propanolol': propanolol,
        'dialisis': dialisis,
        'portal_vein_thrombosis': portal_vein_thrombosis,
        'ascitis': ascitis,
        'hepatocellular_carcinoma': hepatocellular_carcinoma,
        'albumin': albumin,
        'total_bilirrubin': total_bilirrubin,
        'direct_bilirrubina': direct_bilirrubina,
        'inr': inr,
        'creatinine': creatinine,
        'platelets': platelets,
        'ast': ast,
        'alt': alt,
        'hemoglobin': hemoglobin,
        'hematocrit': hematocrit,
        'leucocytes': leucocytes,
        'sodium': sodium,
        'potassium': potassium,
        'varices': varices,
        'red_wale_marks': red_wale_marks,
        'rupture_point': rupture_point,
        'active_bleeding': active_bleeding,
        'therapy': therapy,
        'terlipressin_dose': terlipressin_dose,
        'time-to-endoscophy_hours': time_to_endoscophy_hours,
        'rebleeding': rebleeding
    }
    
    df = pd.DataFrame([input_data])
    
    # Preprocess the input data
    processed_data = preprocessor.transform(df)
    
    # Make predictions
    prediction = model.predict(processed_data)
    probability = model.predict_proba(processed_data)[:, 1]
    
    return prediction[0], probability[0]

# Define Gradio interface
inputs = [
    gr.Slider(minimum=18, maximum=100, label='Age', step=1),
    gr.Dropdown(choices=['male', 'female'], label='Sex'),
    gr.Dropdown(choices=['white', 'black', 'asian', 'other'], label='Race'),
    gr.Dropdown(choices=['alcohol', 'hcv', 'alcohol+hcv', 'other'], label='Etiology Cirrhosis'),
    gr.Dropdown(choices=['yes', 'no'], label='Hepatorenal Syndrome'),
    gr.Dropdown(choices=['yes', 'no'], label='Omeprazole'),
    gr.Dropdown(choices=['yes', 'no'], label='Spironolactone'),
    gr.Dropdown(choices=['yes', 'no'], label='Furosemide'),
    gr.Dropdown(choices=['yes', 'no'], label='Propanolol'),
    gr.Dropdown(choices=['yes', 'no'], label='Dialysis'),
    gr.Dropdown(choices=['yes', 'no'], label='Portal Vein Thrombosis'),
    gr.Dropdown(choices=['yes', 'no'], label='Ascites'),
    gr.Dropdown(choices=['yes', 'no'], label='Hepatocellular Carcinoma'),
    gr.Slider(minimum=1, maximum=5, label='Albumin', step=0.1),
    gr.Slider(minimum=0.1, maximum=30, label='Total Bilirubin', step=0.1),
    gr.Slider(minimum=0.1, maximum=10, label='Direct Bilirubin', step=0.1),
    gr.Slider(minimum=0.5, maximum=5, label='INR', step=0.1),
    gr.Slider(minimum=0.1, maximum=10, label='Creatinine', step=0.1),
    gr.Slider(minimum=10, maximum=500, label='Platelets', step=1),
    gr.Slider(minimum=10, maximum=500, label='AST', step=1),
    gr.Slider(minimum=10, maximum=500, label='ALT', step=1),
    gr.Slider(minimum=5, maximum=20, label='Hemoglobin', step=0.1),
    gr.Slider(minimum=15, maximum=60, label='Hematocrit', step=0.1),
    gr.Slider(minimum=1, maximum=50, label='Leukocytes', step=0.1),
    gr.Slider(minimum=120, maximum=160, label='Sodium', step=1),
    gr.Slider(minimum=2, maximum=6, label='Potassium', step=0.1),
    gr.Dropdown(choices=['yes', 'no'], label='Varices'),
    gr.Dropdown(choices=['yes', 'no'], label='Red Wale Marks'),
    gr.Dropdown(choices=['yes', 'no'], label='Rupture Point'),
    gr.Dropdown(choices=['yes', 'no'], label='Active Bleeding'),
    gr.Dropdown(choices=['Banding', 'Sclerotherapy', 'No therapy'], label='Therapy'),
    gr.Slider(minimum=0, maximum=20, label='Terlipressin Dose', step=1),
    gr.Slider(minimum=0, maximum=48, label='Time to Endoscopy Hours', step=1),
    gr.Dropdown(choices=['yes', 'no'], label='Rebleeding')
]

outputs = [
    gr.Textbox(label='Prediction'),
    gr.Textbox(label='Probability')
]

demo = gr.Interface(fn=predict, inputs=inputs, outputs=outputs, live=False)

demo.launch()