import pickle
import gradio as gr

# Load the trained model
with open("model.pkl", "rb") as f:
    model = pickle.load(f)

# Mapping
diagnoses = {
    0: "Negative",
    1: "Hypothyroid",
    2: "Hyperthyroid"
}

# Preprocess inputs
def preprocess_inputs(age, sex, on_thyroxine, query_on_thyroxine, on_antithyroid_meds, sick, pregnant,
                      thyroid_surgery, I131_treatment, query_hypothyroid, query_hyperthyroid, lithium,
                      goitre, tumor, hypopituitary, psych, TSH, T3, TT4, T4U, FTI):
    
    binary_map = {'Yes': 1, 'No': 0, '': None}
    
    on_thyroxine = binary_map[on_thyroxine]
    query_on_thyroxine = binary_map[query_on_thyroxine]
    on_antithyroid_meds = binary_map[on_antithyroid_meds]
    sick = binary_map[sick]
    pregnant = binary_map[pregnant]
    thyroid_surgery = binary_map[thyroid_surgery]
    I131_treatment = binary_map[I131_treatment]
    query_hypothyroid = binary_map[query_hypothyroid]
    query_hyperthyroid = binary_map[query_hyperthyroid]
    lithium = binary_map[lithium]
    goitre = binary_map[goitre]
    tumor = binary_map[tumor]
    hypopituitary = binary_map[hypopituitary]
    psych = binary_map[psych]
    
    sex = 1 if sex == 'F' else 0 if sex == 'M' else None

    return [age, sex, on_thyroxine, query_on_thyroxine, on_antithyroid_meds, sick, pregnant,
            thyroid_surgery, I131_treatment, query_hypothyroid, query_hyperthyroid, lithium,
            goitre, tumor, hypopituitary, psych, TSH, T3, TT4, T4U, FTI]

# Predict function
def predict_diagnosis(age, sex, on_thyroxine, query_on_thyroxine, on_antithyroid_meds, sick, pregnant,
                      thyroid_surgery, I131_treatment, query_hypothyroid, query_hyperthyroid, lithium,
                      goitre, tumor, hypopituitary, psych, TSH, T3, TT4, T4U, FTI):
    
    inputs = preprocess_inputs(age, sex, on_thyroxine, query_on_thyroxine, on_antithyroid_meds, sick,
                               pregnant, thyroid_surgery, I131_treatment, query_hypothyroid,
                               query_hyperthyroid, lithium, goitre, tumor, hypopituitary, psych,
                               TSH, T3, TT4, T4U, FTI)
    
    output = model.predict([inputs])[0]
    return diagnoses.get(output, "Unknown")

# Define Gradio interface
demo = gr.Interface(
    fn=predict_diagnosis,
    inputs=[
        gr.Number(label="Age"),
        gr.Radio(["M", "F"], label="Sex"),
        gr.Radio(["Yes", "No"], label="On Thyroxine"),
        gr.Radio(["Yes", "No"], label="Query on Thyroxine"),
        gr.Radio(["Yes", "No"], label="On Antithyroid Meds"),
        gr.Radio(["Yes", "No"], label="Sick"),
        gr.Radio(["Yes", "No"], label="Pregnant"),
        gr.Radio(["Yes", "No"], label="Thyroid Surgery"),
        gr.Radio(["Yes", "No"], label="I131 Treatment"),
        gr.Radio(["Yes", "No"], label="Query Hypothyroid"),
        gr.Radio(["Yes", "No"], label="Query Hyperthyroid"),
        gr.Radio(["Yes", "No"], label="Lithium"),
        gr.Radio(["Yes", "No"], label="Goitre"),
        gr.Radio(["Yes", "No"], label="Tumor"),
        gr.Radio(["Yes", "No"], label="Hypopituitary"),
        gr.Radio(["Yes", "No"], label="Psych"),
        gr.Number(label="TSH"),
        gr.Number(label="T3"),
        gr.Number(label="TT4"),
        gr.Number(label="T4U"),
        gr.Number(label="FTI"),
    ],
    outputs="text",
    title="Thyroid Diagnosis Predictor",
    description="Predict whether a patient has Hypothyroid, Hyperthyroid or is Negative using medical features."
)


# Launch app
if __name__ == "__main__":
    demo.launch()