from annotated_text import annotated_text, parameters, annotation
from razdel import tokenize
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
import streamlit as st
import torch

# add the caching decorator and use custom text for spinner
@st.cache_resource(show_spinner = "Loading the model...")

def label_text(text):
	if text != "":
		tokenizer = AutoTokenizer.from_pretrained("yeshpanovrustem/xlm-roberta-large-ner-kazakh")
		model = AutoModelForTokenClassification.from_pretrained("yeshpanovrustem/xlm-roberta-large-ner-kazakh")
		nlp = pipeline("ner", model = model, tokenizer = tokenizer)

		labels_dict = {0: 'O', 
				   1: 'B-ADAGE', 
				   2: 'I-ADAGE', 
				   3: 'B-ART', 
				   4: 'I-ART', 
				   5: 'B-CARDINAL', 
				   6: 'I-CARDINAL', 
				   7: 'B-CONTACT', 
				   8: 'I-CONTACT', 
				   9: 'B-DATE', 
				   10: 'I-DATE', 
				   11: 'B-DISEASE', 
				   12: 'I-DISEASE', 
				   13: 'B-EVENT', 
				   14: 'I-EVENT', 
				   15: 'B-FACILITY', 
				   16: 'I-FACILITY', 
				   17: 'B-GPE', 
				   18: 'I-GPE', 
				   19: 'B-LANGUAGE', 
				   20: 'I-LANGUAGE', 
				   21: 'B-LAW', 
				   22: 'I-LAW', 
				   23: 'B-LOCATION', 
				   24: 'I-LOCATION', 
				   25: 'B-MISCELLANEOUS', 
				   26: 'I-MISCELLANEOUS', 
				   27: 'B-MONEY', 
				   28: 'I-MONEY', 
				   29: 'B-NON_HUMAN', 
				   30: 'I-NON_HUMAN', 
				   31: 'B-NORP', 
				   32: 'I-NORP', 
				   33: 'B-ORDINAL', 
				   34: 'I-ORDINAL', 
				   35: 'B-ORGANISATION', 
				   36: 'I-ORGANISATION', 
				   37: 'B-PERSON', 
				   38: 'I-PERSON', 
				   39: 'B-PERCENTAGE', 
				   40: 'I-PERCENTAGE', 
				   41: 'B-POSITION', 
				   42: 'I-POSITION', 
				   43: 'B-PRODUCT', 
				   44: 'I-PRODUCT', 
				   45: 'B-PROJECT', 
				   46: 'I-PROJECT', 
				   47: 'B-QUANTITY', 
				   48: 'I-QUANTITY', 
				   49: 'B-TIME', 
				   50: 'I-TIME'}

		single_sentence_tokens = [_.text for _ in list(tokenize(text))]
		tokenized_input = tokenizer(single_sentence_tokens, is_split_into_words = True, return_tensors = "pt")
		tokens = tokenized_input.tokens()
		output = model(**tokenized_input).logits
		predictions = torch.argmax(output, dim = 2)

		# convert label IDs to label names
		word_ids = tokenized_input.word_ids(batch_index = 0)
		previous_word_id = None
		labels = []
		for token, word_id, prediction in zip(tokens, word_ids, predictions[0].numpy()):
		#         # Special tokens have a word id that is None. We set the label to -100 so they are
		#         # automatically ignored in the loss function.
			if word_id is None or word_id == previous_word_id:
				continue
			elif word_id != previous_word_id:
				labels.append(labels_dict[prediction])
			previous_word_id = word_id
		assert len(single_sentence_tokens) == len(labels), "Mismatch between input token and label sizes!"

		sentence_tokens = []
		sentence_labels = []

		token_list = []
		label_list = []

		previous_token = ""
		previous_label = ""

		for token, label in zip(single_sentence_tokens, labels):
		    current_token = token
		    current_label = label
		    
		    # starting loop
		    if previous_label == "":
		        previous_token = current_token
		        previous_label = current_label
		        
		    # collecting compound named entities
		    elif (previous_label.startswith("B-")) and (current_label.startswith("I-")):
		        token_list.append(previous_token)
		        label_list.append(previous_label)
		    elif (previous_label.startswith("I-")) and (current_label.startswith("I-")):
		        token_list.append(previous_token)
		        label_list.append(previous_label)
		    elif (previous_label.startswith("I-")) and (not current_label.startswith("I-")):
		        token_list.append(previous_token)
		        label_list.append(previous_label)
		        sentence_tokens.append(token_list)
		        sentence_labels.append(label_list)
		        token_list = []
		        label_list = []
		    # collecting single named entities:
		    elif (not previous_label.startswith("I-")) and (not current_label.startswith("I-")):
		        token_list.append(previous_token)
		        label_list.append(previous_label)
		        sentence_tokens.append(token_list)
		        sentence_labels.append(label_list)
		        token_list = []
		        label_list = []
		    previous_token = current_token
		    previous_label = current_label
		token_list.append(previous_token)
		label_list.append(previous_label)
		sentence_tokens.append(token_list)
		sentence_labels.append(label_list)

		output = []
		for sentence_token, sentence_label in zip(sentence_tokens, sentence_labels):
		    if len(sentence_label[0]) > 1:
		        if len(sentence_label) > 1:
		            output.append((" ".join(sentence_token), sentence_label[0].split("-")[1]))
		        else:
		            output.append((sentence_token[0], sentence_label[0].split("-")[1]))
		    else:
		        # output.append((sentence_token[0], sentence_label[0]))
		        output.append(sentence_token[0])

		modified_output = []
		for element in output:
			if not isinstance(element, tuple):
				if element.isalnum():
					modified_output.append(' ' + element + ' ')
				else:
					modified_output.append(' ' + element + ' ')
			else:
				tuple_first = f" {element[0]} "
				tuple_second = element[1]
				new_tuple = (tuple_first, tuple_second)
				modified_output.append(new_tuple)
	else:
		return st.markdown("<p id = 'warning'>EN: PLEASE INSERT YOUR TEXT<br>ҚАЗ: МӘТІНДІ ОРНАЛАСТЫРЫҢЫЗ<br>РУС: ВСТАВЬТЕ ТЕКСТ</p>", unsafe_allow_html = True)
	return annotated_text(modified_output)

#########################
#### CREATE SIDEBAR #####
#########################

with open("style.css") as f:
	css = f.read()

languages = ['🇬🇧', '🇰🇿', '🇷🇺']

language = st.sidebar.radio("", options = languages, horizontal = True)

st.sidebar.markdown(f'<style>{css}</style>', unsafe_allow_html = True)

if language == languages[0]:
	st.sidebar.markdown("<h1>Kazakh NER</h1>", unsafe_allow_html = True)
	st.sidebar.markdown("<h2>Named entity classes</h2>", unsafe_allow_html = True)

	with st.sidebar.expander("ADAGE"): st.write("Well-known Kazakh proverbs and sayings")
	with st.sidebar.expander("ART"): st.write("Titles of books, songs, television programmes, etc.")
	with st.sidebar.expander("CARDINAL"): st.write("Cardinal numbers, including whole numbers, fractions, and decimals")
	with st.sidebar.expander("CONTACT"): st.write("Addresses, emails, phone numbers, URLs")
	with st.sidebar.expander("DATE"): st.write("Dates or periods of 24 hours or more")
	with st.sidebar.expander("DISEASE"): st.write("Diseases or medical conditions")
	with st.sidebar.expander("EVENT"): st.write("Named events and phenomena")
	with st.sidebar.expander("FACILITY"): st.write("Names of man-made structures")
	with st.sidebar.expander("GPE"): st.write("Names of geopolitical entities")
	with st.sidebar.expander("LANGUAGE"): st.write("Named languages")
	with st.sidebar.expander("LAW"): st.write("Named legal documents")
	with st.sidebar.expander("LOCATION"): st.write("Names of geographical locations other than GPEs")
	with st.sidebar.expander("MISCELLANEOUS"): st.write("Entities of interest but hard to assign a proper tag to")
	with st.sidebar.expander("MONEY"): st.write("Monetary values")
	with st.sidebar.expander("NON_HUMAN"): st.write("Names of pets, animals or non-human creatures")
	with st.sidebar.expander("NORP"): st.write("Adjectival forms of GPE and LOCATION; named religions, etc.")
	with st.sidebar.expander("ORDINAL"): st.write("Ordinal numbers, including adverbials")
	with st.sidebar.expander("ORGANISATION"): st.write("Names of companies, government agencies, etc.")
	with st.sidebar.expander("PERCENTAGE"): st.write("Percentages")
	with st.sidebar.expander("PERSON"): st.write("Names of persons")
	with st.sidebar.expander("POSITION"): st.write("Names of posts and job titles")
	with st.sidebar.expander("PRODUCT"): st.write("Names of products")
	with st.sidebar.expander("PROJECT"): st.write("Names of projects, policies, plans, etc.")
	with st.sidebar.expander("QUANTITY"): st.write("Length, distance, etc. measurements")
	with st.sidebar.expander("TIME"): st.write("Times of day and time duration less than 24 hours")

elif language == languages[1]:
	st.sidebar.markdown("<h1>Атаулы мәндерді анықтау</h1>", unsafe_allow_html = True)
	st.sidebar.markdown("<h2>Атаулы мән түрлері</h2>", unsafe_allow_html = True)

	with st.sidebar.expander("ADAGE >> НАҚЫЛ СӨЗ"): st.write("Қазақ мақал-мәтелдері")
	with st.sidebar.expander("ART >> ТУЫНДЫ"): st.write("Өнер туындылары, теледидар бағдарламалары, ғылыми мәтіндер атаулары")
	with st.sidebar.expander("CARDINAL >> ЕСЕПТІК САН"): st.write("Бүтін сандар мен (ондық) бөлшектер")
	with st.sidebar.expander("CONTACT >> БАЙЛАНЫС ДЕРЕКТЕРІ"): st.write("Мекенжайлар, телефон нөмірлері, сілтемелер және электрондық пошта адрестері")
	with st.sidebar.expander("DATE >> ДАТА"): st.write("Даталар, күндер, апталар, айлар, атауы бар белгілі кезеңдер, мезгілдер, жылдар")
	with st.sidebar.expander("DISEASE >> АУРУ"): st.write("Ауру немесе бұзылу атаулары")
	with st.sidebar.expander("EVENT >> ОҚИҒА"): st.write("Атауы бар оқиғалар мен табиғи құбылыстар")
	with st.sidebar.expander("FACILITY >> ҚҰРЫЛЫС"): st.write("Адам қолымен жасалынған құрылымдардың атаулары")
	with st.sidebar.expander("GPE >> ГӘС"): st.write("Географиялық әкімшілік бірліктерінің атаулары")
	with st.sidebar.expander("LANGUAGE >> ТІЛ"): st.write("Тіл атаулары")
	with st.sidebar.expander("LAW >> ЗАҢ"): st.write("Заңдастырылған құжаттар атаулары")
	with st.sidebar.expander("LOCATION >> МЕКЕН"): st.write("GPE атаулы мәніне жатпайтын географиялық нысандардың атаулары")
	with st.sidebar.expander("MISCELLANEOUS >> АРАЛАС"): st.write("Атаулы мәнді анықтауға қиыншылық тудыратын сөздер, фразалар")
	with st.sidebar.expander("MONEY >> АҚША БІРЛІГІ"): st.write("Құрамында сан және ақша бірлігі айқын берілген фразалар.")
	with st.sidebar.expander("NON_HUMAN >> АДАМ ЕМЕС"): st.write("Жануарлардың және ойдан шығарылған жануар-кейіпкерлердің лақап аттары")
	with st.sidebar.expander("NORP >> ҰДСҚ"): st.write("GPE және LOCATION атаулы мәндеріне қатыстылықты білдіретін сөздер")
	with st.sidebar.expander("ORDINAL >> РЕТТІК САН"): st.write("Реттік сандар және реттік сандардан құрылған сын есімдер мен қыстырма сөздер")
	with st.sidebar.expander("ORGANISATION >> ҰЙЫМ"): st.write("Ұйым атаулары")
	with st.sidebar.expander("PERCENTAGE >> ПАЙЫЗ"): st.write("Құрамында сан және пайыз, процент сөздері немесе % белгісі бар фразалар")
	with st.sidebar.expander("PERSON >> АДАМ"): st.write("Адам есімі, фамилиясы, әкесінің аты")
	with st.sidebar.expander("POSITION >> ОРЫН"): st.write("Адамның белгілі бір ұйымның иерархиялық жүйесіндегі орны немесе атқаратын қызметі")
	with st.sidebar.expander("PRODUCT >> ӨНІМ"): st.write("Өнім атаулары")
	with st.sidebar.expander("PROJECT >> ЖОБА"): st.write("Жобалар, жоспарлар, бастамалар, стратегиялардың атаулары")
	with st.sidebar.expander("QUANTITY >> ӨЛШЕМ"): st.write("Құрамында сан және өлшем бірлігі айқын берілген фразалар")
	with st.sidebar.expander("TIME >> УАҚЫТ"): st.write("Бір тәуліктен қысқа кезеңдер")

elif language == languages[2]:
	st.sidebar.markdown("<h1>Распознавание именованных сущностей</h1>", unsafe_allow_html = True)
	st.sidebar.markdown("<h2>Именованные сущности</h2>", unsafe_allow_html = True)

	with st.sidebar.expander("ADAGE >> ПОСЛОВИЦА"): st.write("Известные казахские пословицы и поговорки")
	with st.sidebar.expander("ART >> ИСКУССТВО"): st.write("Названия книг, песен, телевизионных и радиопрограмм и др.")
	with st.sidebar.expander("CARDINAL >> КОЛИЧ. ЧИСЛ."): st.write("Количественные числительные (целые числа, дроби и десятичные дроби)")
	with st.sidebar.expander("CONTACT >> КОНТ. ДАННЫЕ"): st.write("Адреса, номера телефонов, URL-адреса и др")
	with st.sidebar.expander("DATE >> ДАТА"): st.write("Даты или периоды, продолжительность которых составляет 24 часа и более")
	with st.sidebar.expander("DISEASE >> БОЛЕЗНЬ"): st.write("Названия болезней, заболеваний, различных расстройств и недомоганий")
	with st.sidebar.expander("EVENT >> СОБЫТИЕ"): st.write("Названия мероприятий и природных явлений")
	with st.sidebar.expander("FACILITY >> СООРУЖЕНИЕ"): st.write("Названия зданий и сооружений, созданных руками человека")
	with st.sidebar.expander("GPE >> ГПС"): st.write("Названия геополитических объектов")
	with st.sidebar.expander("LANGUAGE >> ЯЗЫК"): st.write("Наименования и самоназвания языков")
	with st.sidebar.expander("LAW >> ЗАКОН"): st.write("Наименования юридических документов")
	with st.sidebar.expander("LOCATION >> МЕСТО"): st.write("Названия географических местоположений, отличных от GPE")
	with st.sidebar.expander("MISCELLANEOUS >> РАЗНОЕ"): st.write("Сущности, представляющие сложность при классификации")
	with st.sidebar.expander("MONEY >> ДЕНЬГИ"): st.write("Номинал и наименования денежных единиц")
	with st.sidebar.expander("NON_HUMAN >> НЕЧЕЛОВЕК"): st.write("Имена и клички домашних животных, зверей, вымышленных созданий")
	with st.sidebar.expander("NORP >> НОРП"): st.write("Прилагательные, образованные от сущностей GPE и LOCATION")
	with st.sidebar.expander("ORDINAL >> ПОРЯД. ЧИСЛ."): st.write("Порядковые числительные")
	with st.sidebar.expander("ORGANISATION >> ОРГАНИЗАЦИЯ"): st.write("Наименования компаний, правительственных учреждений и т. д.")
	with st.sidebar.expander("PERCENTAGE >> ПРОЦЕНТ"): st.write("Обозначение процентных соотношений")
	with st.sidebar.expander("PERSON >> ЧЕЛОВЕК"): st.write("Имена людей")
	with st.sidebar.expander("POSITION >> ДОЛЖНОСТЬ"): st.write("Наименования должностей")
	with st.sidebar.expander("PRODUCT >> ПРОДУКТ"): st.write("Названия продуктов/изделий различных видов производства")
	with st.sidebar.expander("PROJECT >> ПРОЕКТ"): st.write("Наименования проектов, планов, стратегий, кампаний, референдумов и др")
	with st.sidebar.expander("QUANTITY >> ВЕЛИЧИНА"): st.write("Наименования единиц измерения (длина, расстояние, вес, напряжение, др.)")
	with st.sidebar.expander("TIME >> ВРЕМЯ"): st.write("Обозначение времени, продолжительность которого менее 24 часов")


######################
#### CREATE FORM #####
######################

if language == languages[0]:
	text_field = st.text_area('Insert your text here')
	submit = st.button("Submit", type = "primary")
	st.markdown('Press **Submit** to have your text labelled')
elif language == languages[1]:
	text_field = st.text_area('Мәтінді орналастырыңыз')
	submit = st.button("Анықтау", type = "primary")
	st.markdown('Мәтінді белгілеу үшін **Анықтау** батырмасын басыңыз')
elif language == languages[2]:
	text_field = st.text_area('Вставьте текст')
	submit = st.button("Распознать", type = "primary")
	st.markdown('Нажмите кнопку **Распознать**, чтобы разметить текст')

if submit:
	label_text(text_field)