import gradio as gr
import pickle
import pandas as pd 
import numpy as np
import lightgbm as lgb
from autogluon.tabular import TabularPredictor

loaded_model = pickle.load(open('Autogluon/models/XGBoost_BAG_L1/model.pkl', 'rb'))
# pred_proba_radio = loaded_model.predict_proba(X_test_radio)
# pred_proba_radio


def relapse(age, pathology, B_symptoms):
    
    X_test_radio = pd.DataFrame.from_dict(
        {
            "Age": [age],
            "pathology_Nodular Sclerosis cHL": [1 if pathology == 'Nodular' else 0],
            "pathology_others": [1 if pathology == 'Others' else 0],
            "B_symptoms_Yes" : [1 if B_symptoms else 0],
            'radio_Yes': [1]
        }
    )
    
    X_test_no_radio = pd.DataFrame.from_dict(
        {
            "Age": [age],
            "pathology_Nodular Sclerosis cHL": [1 if pathology == 'Nodular' else 0],
            "pathology_others": [1 if pathology == 'Others' else 0],
            "B_symptoms_Yes" : [1 if B_symptoms else 0],
            'radio_Yes': [0]
        }
    )    

    pred_proba_radio = loaded_model.predict_proba(X_test_radio)
    
    pred_proba_radio =  round(np.ndarray.item(pred_proba_radio),2)

    pred_radio = loaded_model.predict(X_test_radio)
    
    pred_proba_no_radio = loaded_model.predict_proba(X_test_no_radio)  
    
    pred_proba_no_radio =  round(np.ndarray.item(pred_proba_no_radio),2)

    pred_no_radio = loaded_model.predict(X_test_no_radio)
    
    return {"Radio": pred_proba_radio, "No Radio": pred_proba_no_radio}

iface = gr.Interface(
    
    title = 'Should we omit radiotherapy?', 
                     description = 'This model predicts relapse according to risk factors.', 
                    fn=relapse, 
                     inputs= [
                         gr.Number(label = 'Age', show_label = True),
        gr.Radio( choices = ['Mixed cellularity', 'Nodular', 'Others'], label = 'Pathology', show_label = True),
                        gr.Checkbox(label = 'B_symptoms', show_label = True)],
    outputs = gr.Label(label = 'Has a higher probability of relapse', show_label = True)
    ,live = True,
    interpretation="default",
)

iface.launch()