import streamlit as st
import pandas as pd
import numpy as np
from unidecode import unidecode
import tensorflow as tf 
import cloudpickle
from transformers import DistilBertTokenizerFast
import os
from matplotlib import pyplot as plt
from PIL import Image


with open(os.path.join("models", "toxic_comment_preprocessor_classnames.bin"), "rb") as model_file_obj:
        text_preprocessor, class_names = cloudpickle.load(model_file_obj)    
interpreter = tf.lite.Interpreter(model_path=os.path.join("models", "toxic_comment_classifier_hf_distilbert.tflite"))


def inference(text):
    text = text_preprocessor.preprocess(pd.Series(text))[0]
    model_checkpoint = "distilbert-base-uncased"
    tokenizer = DistilBertTokenizerFast.from_pretrained(model_checkpoint)
    tokens = tokenizer(text, max_length=512, padding="max_length", truncation=True, return_tensors="tf")
    
    # tflite model inference  
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()[0]
    attention_mask, input_ids = tokens['attention_mask'], tokens['input_ids']
    interpreter.set_tensor(input_details[0]["index"], attention_mask)
    interpreter.set_tensor(input_details[1]["index"], input_ids)
    interpreter.invoke()
    tflite_pred = interpreter.get_tensor(output_details["index"])[0]
    result_df = pd.DataFrame({'class': class_names, 'prob': tflite_pred})
    result_df.sort_values(by='prob', ascending=True, inplace=True)
    return result_df


def display_image(df):
    fig, ax = plt.subplots(figsize=(2, 1.8))
    df.plot(x='class', y='prob', kind='barh', ax=ax, color='black', ylabel='')
    ax.tick_params(axis='both', which='major', labelsize=8.5)
    ax.get_legend().remove()
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['bottom'].set_visible(False)
    ax.spines['left'].set_visible(False)
    ax.get_xaxis().set_ticks([])
    plt.rcParams["figure.autolayout"] = True
    plt.xlim(0, 1)
    for n, i in enumerate([*df['prob']]):
        plt.text(i+0.015, n-0.15, f'{str(np.round(i, 3))}   ', fontsize=7.5)

    fig.savefig("prediction.png", bbox_inches='tight', dpi=100)
    image = Image.open('prediction.png')
    st.write('')
    st.image(image, output_format="PNG", caption="Prediction")
    
############## ENTRY POINT START #######################
def main():
    st.title("Toxic Comment Classifier")
    comment_txt = st.text_area("Enter a comment:", "", height=100)
    if st.button("Submit"):
        df = inference(comment_txt)
        display_image(df)

############## ENTRY POINT END #######################

if __name__ == "__main__":
    main()