import gradio as gr
import pickle
import subprocess
import sys


inputs_fields = ['Age',
                 'Gender',
                 'Smoking',
                 'Hx Smoking',
                 'Hx Radiothreapy',
                 'Thyroid Function',
                 'Physical Examination',
                 'Adenopathy',
                 'Pathology',
                 'Focality',
                 'Risk',
                 'T',
                 'N',
                 'M',
                 'Stage',
                 'Response'
                 ]

inputs_for_categorical_fields_values = {
    'Gender' : ['F', 'M'],
    'Smoking': ['No', 'Yes'],
    'Hx Smoking (Smoking History)' :['No', 'Yes'],
    'Hx Radiothreapy (Radiotherapy History)':['No', 'Yes'],
    'Thyroid Function':['Euthyroid', 'Clinical Hyperthyroidism', 'Clinical Hypothyroidism'
        , 'Subclinical Hyperthyroidism', 'Subclinical Hypothyroidism'],
    'Physical Examination':['Single nodular goiter-left', 'Multinodular goiter'
        , 'Single nodular goiter-right', 'Normal', 'Diffuse goiter'],
    'Adenopathy': ['No', 'Right', 'Extensive', 'Left', 'Bilateral', 'Posterior'],
    'Pathology':['Micropapillary', 'Papillary', 'Follicular', 'Hurthel cell'],
    'Focality':['Uni-Focal', 'Multi-Focal'],
    'Risk':['Low', 'Intermediate', 'High'],
    'Tumor':['T1a (tumor that is 1 cm or smaller)', 'T1b (tumor between 1cm and 2cm)', 'T2 (tumor between 2cm and 4cm)',
         'T3a (tumor larger than 4 cm)', 'T3b (tumor that has grown outside the thyroid)',
         'T4a (tumor that has invaded nearby structures)', 'T4b (tumor that has invaded nearby structures)'],
    'Lymph Nodes':['N0 (no evidence of regional lymph node metastasis)',
         'N1b (regional lymph node metastasis in the central  of the neck)',
         'N1a (regional lymph node metastasis in the lateral of the neck)'],
    'Cancer Metastasis':['M0 (no evidence of distant metastasis)', 'M1 (the presence of distant metastasis)'],
    'Stage':['I', 'II', 'IVB', 'III', 'IVA'],
    'Response':['Indeterminate', 'Excellent', 'Structural Incomplete', 'Biochemical Incomplete']
}

expected_inputs = ['Age',
                   'Gender_F',
                   'Gender_M',
                   'Smoking_No',
                   'Smoking_Yes',
                   'Hx Smoking_No',
                   'Hx Smoking_Yes',
                   'Hx Radiothreapy_No',
                   'Hx Radiothreapy_Yes',
                   'Thyroid Function_Clinical Hyperthyroidism',
                   'Thyroid Function_Clinical Hypothyroidism',
                   'Thyroid Function_Euthyroid',
                   'Thyroid Function_Subclinical Hyperthyroidism',
                   'Thyroid Function_Subclinical Hypothyroidism',
                   'Physical Examination_Diffuse goiter',
                   'Physical Examination_Multinodular goiter',
                   'Physical Examination_Normal',
                   'Physical Examination_Single nodular goiter-left',
                   'Physical Examination_Single nodular goiter-right',
                   'Adenopathy_Bilateral',
                   'Adenopathy_Extensive',
                   'Adenopathy_Left',
                   'Adenopathy_No',
                   'Adenopathy_Posterior',
                   'Adenopathy_Right',
                   'Pathology_Follicular',
                   'Pathology_Hurthel cell',
                   'Pathology_Micropapillary',
                   'Pathology_Papillary',
                   'Focality_Multi-Focal',
                   'Focality_Uni-Focal',
                   'Risk_High',
                   'Risk_Intermediate',
                   'Risk_Low',
                   'T_T1a',
                   'T_T1b',
                   'T_T2',
                   'T_T3a',
                   'T_T3b',
                   'T_T4a',
                   'T_T4b',
                   'N_N0',
                   'N_N1a',
                   'N_N1b',
                   'M_M0',
                   'M_M1',
                   'Stage_I',
                   'Stage_II',
                   'Stage_III',
                   'Stage_IVA',
                   'Stage_IVB',
                   'Response_Biochemical Incomplete',
                   'Response_Excellent',
                   'Response_Indeterminate',
                   'Response_Structural Incomplete']

def normalize_age(user_age, age_min=15, age_max=82):
    user_age = int(user_age)
    assert age_min <= user_age <= age_max, f"Age must be between {age_min} and {age_max}"
    assert user_age >= 0, "Age must be a positive integer"
    # Normalize age using the min and max from training
    normalized_age = (user_age - age_min) / (age_max - age_min)
    return int(normalized_age)

def transform_input_to_expected_format(user_input):
    # Initialize output dictionary with all expected inputs set to 0
    transformed_input = {feature: 0 for feature in expected_inputs}

    for field, value in user_input.items():
        if type(value) == str:
            value = value.split(' (')[0]
        if field == 'Tumor':
            field = 'T'
        if field == 'Lymph Nodes':
            field = 'N'
        if field == 'Cancer Metastasis':
            field = 'M'

        if field == 'Age':
            transformed_input['Age'] = normalize_age(value)
        else:
            key = f"{field}_{value}"
            if key in transformed_input:
                transformed_input[key] = 1

    return transformed_input

def predict_thyroid_cancer(Age, Gender, Smoking, Hx_Smoking, Hx_Radiothreapy,Thyroid_Function, Physical_Examination, Adenopathy, Pathology, Focality, Risk, T, N, M, Stage, Response):
    inputs = {
        'Age': int(Age),
        'Gender': Gender,
        'Smoking': Smoking,
        'Hx Smoking': Hx_Smoking,
        'Hx Radiothreapy': Hx_Radiothreapy,
        'Thyroid Function': Thyroid_Function,
        'Physical Examination': Physical_Examination,
        'Adenopathy': Adenopathy,
        'Pathology': Pathology,
        'Focality': Focality,
        'Risk': Risk,
        'T': T,
        'N': N,
        'M': M,
        'Stage': Stage,
        'Response': Response
    }
    with open('random_forest_model.pkl', 'rb') as model_file:
        model = pickle.load(model_file)

    transformed_input = list(transform_input_to_expected_format(inputs).values())

    # Prediction
    predictions = model.predict([transformed_input])
    risk_level = 'high' if predictions[0] == 1 else 'low'

    # Probabilities
    probabilities = model.predict_proba([transformed_input])[0]
    class_probabilities = dict(zip(model.classes_, probabilities))
    probabilities_str = ", ".join([f"{'low risk' if cls == 0.0 else 'high risk'} at {prob * 100:.2f}%" for cls, prob in class_probabilities.items()])

    return f"Patient has {risk_level} risk of thyroid cancer recurrence.\nProbabilities: {probabilities_str}"


# Intall required packages
packages = ["scikit-learn"]
subprocess.check_call([sys.executable, "-m", "pip", "install"] + packages)


dropdown_inputs = [
    gr.Textbox(label="Age (between 15 and 82)")
]

for field, choices in inputs_for_categorical_fields_values.items():
    dropdown_inputs.append(gr.Dropdown(choices=choices, label=field))

demo = gr.Interface(fn=predict_thyroid_cancer, inputs=dropdown_inputs, outputs="text")
demo.launch(share=True)