import gradio as gr
import numpy as np
import pandas as pd
import joblib
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier

# Sample synthetic dataset creation
def create_dataset():
    np.random.seed(42)
    num_samples = 1000

    age = np.random.randint(18, 90, num_samples)
    severity_score = np.random.uniform(0, 100, num_samples)
    ventilator_support = np.random.choice([0, 1], num_samples)
    length_of_stay = np.random.randint(1, 30, num_samples)

    discharge_destinations = np.random.choice(["Home", "Death", "Nursing Facility", "Rehabilitation"], num_samples)

    data = pd.DataFrame({
        "Age": age,
        "Severity_Score": severity_score,
        "Ventilator_Support": ventilator_support,
        "Length_of_Stay": length_of_stay,
        "Discharge_Destination": discharge_destinations
    })
    
    return data

# Generate dataset
df = create_dataset()

# Train model
X = df.drop(columns=["Discharge_Destination"])
y = df["Discharge_Destination"]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# Save model
joblib.dump(model, "icu_model.pkl")

# Load model
loaded_model = joblib.load("icu_model.pkl")

# Prediction function
def predict_discharge(age, severity_score, ventilator_support, length_of_stay):
    input_data = np.array([[age, severity_score, ventilator_support, length_of_stay]])
    prediction = loaded_model.predict(input_data)[0]
    return prediction

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# ICU Patient Discharge Destination Predictor")
    
    age = gr.Number(label="Age", value=30, minimum=18, maximum=90)
    severity_score = gr.Number(label="Severity Score", value=50, minimum=0, maximum=100)
    ventilator_support = gr.Radio(["No", "Yes"], label="Ventilator Support", value="No", type="index")
    length_of_stay = gr.Number(label="Length of Stay (Days)", value=5, minimum=1, maximum=30)

    output = gr.Textbox(label="Predicted Discharge Destination")

    submit = gr.Button("Predict")
    submit.click(predict_discharge, inputs=[age, severity_score, ventilator_support, length_of_stay], outputs=output)

demo.launch()