import gradio as gr
import joblib
import numpy as np

# Load the model
model_filename = 'gbdt_inv_cd_ratio.joblib'
gbdt = joblib.load(model_filename)

# Define the prediction function
def predict_immune_risk(ca, mg, zn, phos):
    # Prepare input data
    input_data = np.array([[ca, mg, zn, phos]])
    
    # Make predictions
    predicted_class = gbdt.predict(input_data)[0]
    predicted_probabilities = gbdt.predict_proba(input_data)[0]

    # Determine the risk label and probability
    if predicted_class == 1 or predicted_probabilities[1] > 0.5:
        risk_label = "High Immune Risk (Inverted CD4/CD8 Ratio)"
        probability = predicted_probabilities[1] * 100
    else:
        risk_label = "Low Immune Risk (Normal CD4/CD8 Ratio)"
        probability = (1 - predicted_probabilities[1]) * 100  # Using the probability of the other class

    return risk_label, f"{probability:.2f}%"

# Create Gradio interface
iface = gr.Interface(
    fn=predict_immune_risk,
    inputs=[gr.Number(label="Calcium (mg/dL)"),
            gr.Number(label="Magnesium (mg/dL)"),
            gr.Number(label="Zinc (mg/dL)"),
            gr.Number(label="Phosphate (mg/dL)")],
    outputs=[gr.Textbox(label="Immune Risk Prediction"),
             gr.Textbox(label="Predicted Probability")],
    title="HIV Immune Risk Predictor - CD4/CD8 Ratio and Essential Mineral Status",
    description="Enter the levels of Calcium, Magnesium, Zinc, and Phosphate to predict immune risk based on Inversion of CD4/CD8 ratio."
)

# Launch the interface
iface.launch()