print('INFO: import modules')
import gradio as gr
import pickle
from required_classes import *


print('INFO: loading model')
try:
    with open('model_finetuned_clear.pkl', 'rb') as f:
        model = pickle.load(f)
    model.batch_size = 1
    print('INFO: model loaded')
except Exception as e:
    print(f"ERROR: loading models failed with: {str(e)}")

def classify_code(text):
    embed = model._texts2vecs([text])
    probs = model.classifier_code.predict_proba(embed)
    best_n = np.flip(np.argsort(probs, axis=1,)[0])
    preds = {model.classifier_code.classes_[i]: probs[0][i] for i in best_n}
    return preds

def classify_group(text):
    embed = model._texts2vecs([text])
    probs = model.classifier_group.predict_proba(embed)
    best_n = np.flip(np.argsort(probs, axis=1,)[0])
    preds = {model.classifier_group.classes_[i]: probs[0][i] for i in best_n}
    return preds

def classify(text):
    return classify_code(text), classify_group(text)

print('INFO: starting gradio interface')
default_input_text = """Microscopic Evaluation: The specimen is excised into the reticular dermis and demonstrates skin with prominent papillomatosis and hyperkeratosis. The keratinocytes show reactive changes with nuclear enlargement and increased amounts of cytoplasm but no dysplasia is identified. Lymphocytic inflammation is distributed through the base of the lesion. Some pigmentation is present. The lesion approximates the edges of the specimen but appears to be predominantly excised. (GMW/jc) Gross Description: Shave removal and destruction erythematous tender papule with hyperkeratotic scale SCC vs BCC D.. Right medial temple. Final Diagnosis: Pigmented seborrheic keratosis inflamed."""
iface = gr.Interface(
    enable_queue=True,
    title="ICD10-codes classification",
    description="",
    fn=classify,
    inputs=[gr.Textbox(label="Input text", value=default_input_text)],
    outputs=[gr.Label(num_top_classes=4, label="Result class"), gr.Label(num_top_classes=4, label="Result group")],
)
iface.launch()