import streamlit as st
from transformers import DistilBertModel, DistilBertTokenizer
import torch
import matplotlib.image as mpimg

img = mpimg.imread('./460.jpeg')
model_path = './models/pytorch_distilbert.bin'
vocab_path = './models/vocab_distilbert.bin'
device = torch.device('cpu')
MAX_LEN = 512

labels_description = {0: 'Computer Science',
                      1: 'Economics',
                      2: 'Electrical Engineering and Systems Science',
                      3: 'Mathematics',
                      4: 'Physics',
                      5: 'Quantitative Biology',
                      6: 'Quantitative Finance',
                      7: 'Statistics'}


class DistillBERTClass(torch.nn.Module):
    def __init__(self):
        super(DistillBERTClass, self).__init__()
        self.l1 = DistilBertModel.from_pretrained("distilbert-base-cased")
        self.pre_classifier = torch.nn.Linear(768, 768)
        self.dropout = torch.nn.Dropout(0.3)
        self.classifier = torch.nn.Linear(768, 8)

    def forward(self, input_ids, attention_mask):
        output_1 = self.l1(input_ids=input_ids, attention_mask=attention_mask)
        hidden_state = output_1[0]
        pooler = hidden_state[:, 0]
        pooler = self.pre_classifier(pooler)
        pooler = torch.nn.ReLU()(pooler)
        pooler = self.dropout(pooler)
        output = self.classifier(pooler)
        return output


def predict(text, model, human_readable=True):
    model.eval()
    text = " ".join(text.split())
    inputs = tokenizer.encode_plus(
        text,
        None,
        add_special_tokens=True,
        max_length=MAX_LEN,
        pad_to_max_length=True,
        return_token_type_ids=True,
        truncation=True
    )
    ids = torch.tensor(inputs['input_ids'], dtype=torch.long)
    ids = torch.reshape(ids, (1, MAX_LEN))
    mask = torch.tensor(inputs['attention_mask'], dtype=torch.long)
    mask = torch.reshape(mask, (1, MAX_LEN))
    with torch.no_grad():
        outputs = torch.softmax(model(ids, mask), dim=-1)[0].tolist()
    result = []
    for i, v in enumerate(outputs):
        result.append((v, i))
    result.sort(reverse=True)
    pr = 0.0
    index = 0
    answer = []
    while pr < 0.95:
        pr += result[index][0]
        if not human_readable:
            answer.append(result[index][1])
        else:
            answer.append(labels_description[result[index][1]] + " {:.2f}%".format(100 * result[index][0]))
        index += 1

    return answer


tokenizer = DistilBertTokenizer.from_pretrained(vocab_path)
model = torch.load(model_path, map_location=torch.device(device))

st.markdown("### Hi! This is a service for determining the subject of an article.")
st.image(img)
st.markdown("### Just write the title and content in the areas below and click the \"Analyze\" button.")

text1 = st.text_area("Title")

text2 = st.text_area("Summary")

if st.button('Analyse'):
    with st.spinner("Wait..."):
        if text1 or text2:
                pred = predict(text1+"\n"+text2, model.to(device))
                st.success("\n\n".join(pred))
        else:
            st.error(f"You haven't written anything.")