import gradio as gr
import pickle

# Load models using pickle


with open("rf_model.pkl", "rb") as f:
    rf_model = pickle.load(f)

with open("svm_model.pkl", "rb") as f:
    svm_model = pickle.load(f)

# Map for model selection
model_map = {
    "Random Forest": rf_model,
    "SVM": svm_model
}

# Prediction function
def dis_prediction(model_name, sex, pregnant,on_thyroxine, TT4, T3, T4U, FTI, TSH):
    try:
        model = model_map[model_name]

        # Convert input to correct types
        sex = int(sex)
        pregnant = int(pregnant)
        on_thyroxine = int(on_thyroxine)
        TT4 = float(TT4)
        T3 = float(T3)
        T4U = float(T4U)
        FTI = float(FTI)
        TSH = float(TSH)

        # Predict
        result = model.predict([[sex, pregnant, TT4, T3, T4U, FTI, TSH]])
        label_map = {0: "Hyperthyroid", 1: "Hypothyroid", 2: "Negative"}
        return f"Prediction using {model_name}: {label_map.get(result[0], 'Unknown')}"
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio UI
demo = gr.Interface(
    fn=dis_prediction,
    inputs=[
        gr.Dropdown(["SVM",  "Random Forest"], label="Select Model"),
        gr.Radio([0, 1], label="Sex (0: Female, 1: Male)"),
        gr.Radio([0, 1], label="Pregnant (0: No, 1: Yes)"),
        gr.Radio([0, 1], label="On Thyroxine (0: No, 1: Yes)"),
        gr.Number(label="TT4"),
        gr.Number(label="T3"),
        gr.Number(label="T4U"),
        gr.Number(label="FTI"),
        gr.Number(label="TSH"),
    ],
    outputs="text",
    title="Hyperthyroid Prediction (with Pickle Models)",
    description="Choose a model and enter patient data to predict thyroid condition."
)

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