import gradio as gr
import joblib
import numpy as np
import tensorflow as tf

def create_model(input_shape):
    model = tf.keras.models.Sequential([
        tf.keras.layers.Flatten(input_shape=input_shape),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dense(1000, activation='relu'),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(0.8),
        tf.keras.layers.Dense(1000, activation='relu'),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(0.8),
        tf.keras.layers.Dense(1, activation='sigmoid')
    ])
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=3e-6),
                  loss='binary_crossentropy',
                  metrics=['accuracy'])
    return model

model = create_model([18])
model.load_weights('model_2.h5')
scaler = joblib.load('scaler.joblib')

features = [
    'PRNCPTX_HEPATECTOMY RESCJ TOTAL RIGHT LOBECTOMY', 'OPTIME', 'HYPERMED_Yes',
    'PRNCPTX_HEPATECTOMY RESCJ TRISEGMENTECTOMY', 'PRPLATE', 'PRINR',
    'HEP_APPROACH_Laparoscopic', 'PRALKPH', 'HEP_PRINGLE_Yes',
    'HEP_SEC_HISTOLOGIC_Colorectal metastasis', 'HEP_CON_PARTRES', 'SEX_male',
    'ASACLAS', 'HEP_SEC_TUMORSIZE', 'HEP_SEC_NUMTUMORS', 'PRSGOT', 'HEP_VIRAL_None',
    'HEP_NEOTHERAPY_140101_Preoperative systemic chemotherapy'
]

feature_descriptions = {
    'PRNCPTX_HEPATECTOMY RESCJ TOTAL RIGHT LOBECTOMY': "Principal Text: Hepatectomy Resection Total Right Lobectomy",
    'OPTIME': "Operation Time",
    'HYPERMED_Yes': "Hypertension Medication: Yes/No",
    'PRNCPTX_HEPATECTOMY RESCJ TRISEGMENTECTOMY': "Principal Text: Hepatectomy Resection Tri-segmentectomy",
    'PRPLATE': "Preoperative Platelet Count",
    'PRINR': "Preoperative INR Level",
    'HEP_APPROACH_Laparoscopic': "Hepatectomy Approach: Laparoscopic",
    'PRALKPH': "Preoperative Alkaline Phosphatase Level",
    'HEP_PRINGLE_Yes': "Pringle Maneuver Used: Yes/No",
    'HEP_SEC_HISTOLOGIC_Colorectal metastasis': "Secondary Histologic Diagnosis: Colorectal Metastasis",
    'HEP_CON_PARTRES': "Concurrent Partial Resection",
    'SEX_male': "Patient Gender: Male",
    'ASACLAS': "ASA Classification",
    'HEP_SEC_TUMORSIZE': "Secondary Tumor Size",
    'HEP_SEC_NUMTUMORS': "Number of Secondary Tumors",
    'PRSGOT': "Preoperative SGOT Levels",
    'HEP_VIRAL_None': "No Viral Hepatitis",
    'HEP_NEOTHERAPY_140101_Preoperative systemic chemotherapy': "Preoperative Systemic Chemotherapy"
}

categorical_features = [
    'HYPERMED_Yes', 'HEP_APPROACH_Laparoscopic', 'HEP_PRINGLE_Yes',
    'SEX_male', 'HEP_VIRAL_None'
]

def predict_outcome(*args):
    inputs = list(args)
    for i, feature in enumerate(features):
        if feature in categorical_features:
            inputs[i] = 1 if inputs[i] == 'Yes' else 0
    input_data = np.array(inputs).reshape(1, -1)
    inputs_scaled = scaler.transform(input_data)
    prediction = model.predict(inputs_scaled)
    outcome_probability = prediction.flatten()[0]
    prediction_text = f'Probability of Post Hepatectomy Liver Failure: {outcome_probability:.2%}'
    markdown_text = "More details on the variables used in NSQIP can be found in the [NSQIP PUF User Guide 2021]	(https://www.facs.org/media/wd2hlqzv/nsqip_puf_userguide_2021.pdf)."
    return prediction_text, markdown_text

inputs = [gr.Radio(label=feature_descriptions.get(feature), choices=['Yes', 'No']) if feature in categorical_features else gr.Number(label=feature_descriptions.get(feature)) for feature in features]

output_text = gr.Textbox(label="Probability of Liver Failure")
link_info = gr.Markdown()

iface = gr.Interface(
    fn=predict_outcome,
    inputs=[gr.Radio(label=feature_descriptions.get(feature), choices=['Yes', 'No']) if feature in categorical_features else gr.Number(label=feature_descriptions.get(feature)) for feature in features],
    outputs=[output_text, link_info],
    title="NSQIP Post Hepatectomy Liver Failure Calculator"
)