import gradio as gr
import joblib

# Load the trained model
model = joblib.load('trained_random_forest_pipeline.joblib')

def predict(age, sex, race, etiology, hepatorenal_syndrome, omeprazole, spironolactone, furosemide, propanolol, dialysis, portal_vein_thrombosis,
            ascites, hepatocellular_carcinoma, albumin, total_bilirubin, direct_bilirubin, inr, creatinine, platelets, ast, alt, hemoglobin, 
            hematocrit, leukocytes, sodium, potassium, varices, red_wale_marks, rupture_point, active_bleeding, therapy, terlipressin_dose, rebleeding):
    # Convert input data to the format expected by the model, e.g., a list or a DataFrame
    input_data = [age, sex, race, etiology, hepatorenal_syndrome, omeprazole, spironolactone, furosemide, propanolol, dialysis, portal_vein_thrombosis,
                  ascites, hepatocellular_carcinoma, albumin, total_bilirubin, direct_bilirubin, inr, creatinine, platelets, ast, alt, hemoglobin,
                  hematocrit, leukocytes, sodium, potassium, varices, red_wale_marks, rupture_point, active_bleeding, therapy, terlipressin_dose, rebleeding]
    # Assuming the model expects a single sample reshaped as a 2D array
    prediction = model.predict([input_data])
    return str(prediction[0])  # Convert prediction to string if necessary


iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.inputs.Number(label="Age", default=45, minimum=0, maximum=120),
        gr.inputs.Dropdown(label="Sex", choices=["male", "female"]),
        gr.inputs.Dropdown(label="Race", choices=["black", "white", "other"]),
        gr.inputs.Dropdown(label="Etiology", choices=["alcohol", "alcohol+hcv", "hcv", "other"]),
        gr.inputs.Radio(label="Hepatorenal Syndrome", choices=["yes", "no"]),
        gr.inputs.Radio(label="Omeprazole", choices=["yes", "no"]),
        gr.inputs.Radio(label="Spironolactone", choices=["yes", "no"]),
        gr.inputs.Radio(label="Furosemide", choices=["yes", "no"]),
        gr.inputs.Radio(label="Propanolol", choices=["yes", "no"]),
        gr.inputs.Radio(label="Dialysis", choices=["yes", "no"]),
        gr.inputs.Radio(label="Portal Vein Thrombosis", choices=["yes", "no"]),
        gr.inputs.Radio(label="Ascites", choices=["yes", "no"]),
        gr.inputs.Radio(label="Hepatocellular Carcinoma", choices=["yes", "no"]),
        gr.inputs.Number(label="Albumin", default=2.5, step=0.01),
        gr.inputs.Number(label="Total Bilirubin", default=1.2, step=0.01),
        gr.inputs.Number(label="Direct Bilirubin", default=0.5, step=0.01),
        gr.inputs.Number(label="INR", default=1.1, step=0.01),
        gr.inputs.Number(label="Creatinine", default=1.0, step=0.01),
        gr.inputs.Number(label="Platelets", default=150000, step=1),
        gr.inputs.Number(label="AST", default=40, step=1),
        gr.inputs.Number(label="ALT", default=35, step=1),
        gr.inputs.Number(label="Hemoglobin", default=14, step=0.1),
        gr.inputs.Number(label="Hematocrit", default=40, step =0.1),
        gr.inputs.Number(label="Leukocytes", default=5000, step=1),
        gr.inputs.Number(label="Sodium", default=140, step=0.1),
        gr.inputs.Number(label="Potassium", default=4.0, step=0.01),
        gr.inputs.Radio(label="Varices", choices=["yes", "no"]),
        gr.inputs.Radio(label="Red Wale Marks", choices=["yes", "no"]),
        gr.inputs.Radio(label="Rupture Point", choices=["yes", "no"]),
        gr.inputs.Radio(label="Active Bleeding", choices=["yes", "no"]),
        gr.inputs.Dropdown(label="Therapy", choices=["Banding", "no therapy", "Other"]),
        gr.inputs.Number(label="Terlipressin Dose", default=1, minimum=0, maximum=10),
        gr.inputs.Radio(label="Rebleeding", choices=["yes", "no"])
    ],
    outputs="text"
)

if __name__ == "__main__":
    iface.launch()