import streamlit as st
from spacy import displacy
from Model.NER.VLSP2021.Predict_Ner import ViTagger
import re

import torch




device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def process_text(text):
    # Loại bỏ dấu cách thừa và dấu cách ở đầu và cuối văn bản
    processed_text = re.sub(r'\s+', ' ', text.strip())
    return processed_text

def show_ner():

    st.sidebar.title('Datasets')
    dataset = st.sidebar.selectbox("Datasets", ("VLSP2016", "VLSP2021"))
    st.header("NER")
    text = st.text_area("Enter your text for NER:", height=300)
    text = process_text(text)
    if st.button("Process NER"):
        if dataset == "VLSP2021":
                tagger = ViTagger(model_path='Model/NER/VLSP2021/best_model.pt')
                a = text
                b = tagger.extract_entity_doc(a)

                # Danh sách các từ và nhãn NER
                # words_and_labels = [('Dự', 'O'), ('báo', 'O'), ('thời', 'O'), ('tiết', 'O'), ('2 ngày', 'DATETIME-DATERANGE'), ('tới', 'O'), ('Nam   Bộ', 'LOCATION-GPE'), ('có', 'O'), ('mưa', 'O'), ('dông', 'O')]
                words_and_labels = b
                # Tạo danh sách từ
                words = [word for word, _ in words_and_labels]
                # Tạo danh sách thực thể và nhãn cho mỗi từ, loại bỏ nhãn 'O'
                entities = [{'start': sum(len(word) + 1 for word, _ in words_and_labels[:i]),
                             'end': sum(len(word) + 1 for word, _ in words_and_labels[:i + 1]), 'label': label} for
                            i, (word, label)
                            in enumerate(words_and_labels) if label != 'O']
                # print(entities)

                # Render the visualization without color for 'O' labels
                html = displacy.render(
                    {"text": " ".join(words), "ents": entities, "title": None},
                    style="ent",
                    manual=True,
                    options={"colors": {"DATETIME-DATERANGE": "#66c2ff",
                                        "LOCATION-GPE": "#ffcc99",
                                        "O": None,  # Màu cho nhãn 'O'
                                        "QUANTITY-NUM": "#ffdf80",
                                        "EVENT-CUL": "#bfbfbf",
                                        "DATETIME": "#80ff80",
                                        "PERSONTYPE": "#ff80ff",
                                        "PERSON": "#bf80ff",
                                        "QUANTITY-PER": "#80cccc",
                                        "ORGANIZATION": "#ff6666",
                                        "LOCATION-GEO": "#66cc66",
                                        "LOCATION-STRUC": "#cccc66",
                                        "PRODUCT-COM": "#ffff66",
                                        "DATETIME-DATE": "#66cccc",
                                        "QUANTITY-DIM": "#6666ff",
                                        "PRODUCT": "#cc6666",
                                        "QUANTITY": "#6666cc",
                                        "DATETIME-DURATION": "#9966ff",
                                        "QUANTITY-CUR": "#ff9966",
                                        "DATETIME-TIME": "#cdbf93",
                                        "QUANTITY-TEM": "#cc9966",
                                        "DATETIME-TIMERANGE": "#cc8566",
                                        "EVENT-GAMESHOW": "#8c8c5a",
                                        "QUANTITY-AGE": "#70db70",
                                        "QUANTITY-ORD": "#e699ff",
                                        "PRODUCT-LEGAL": "#806699",
                                        "LOCATION": "#993366",
                                        "ORGANIZATION-MED": "#339933",
                                        "URL": "#ff4d4d",
                                        "PHONENUMBER": "#99cc99",
                                        "ORGANIZATION-SPORTS": "#6666ff",
                                        "EVENT-SPORT": "#ffff80",
                                        "SKILL": "#b38f66",
                                        "EVENT-NATURAL": "#ff9966",
                                        "ADDRESS": "#cc9966",
                                        "IP": "#b38f66",
                                        "EMAIL": "#cc8566",
                                        "ORGANIZATION-STOCK": "#666633",
                                        "DATETIME-SET": "#70db70",
                                        "PRODUCT-AWARD": "#e699ff",
                                        "MISCELLANEOUS": "#806699",
                                        "LOCATION-GPE-GEO": "#99ffff"}}
                )
                # print(html)
                st.markdown(html, unsafe_allow_html=True)
        elif dataset == "VLSP2016":
                tagger = ViTagger(model_path='Model/NER/VLSP2016/best_model.pt')
                a = text
                b = tagger.extract_entity_doc(a)

                # Danh sách các từ và nhãn NER
                # words_and_labels = [('Dự', 'O'), ('báo', 'O'), ('thời', 'O'), ('tiết', 'O'), ('2 ngày', 'DATETIME-DATERANGE'), ('tới', 'O'), ('Nam   Bộ', 'LOCATION-GPE'), ('có', 'O'), ('mưa', 'O'), ('dông', 'O')]
                words_and_labels = b
                # Tạo danh sách từ
                words = [word for word, _ in words_and_labels]
                # Tạo danh sách thực thể và nhãn cho mỗi từ, loại bỏ nhãn 'O'
                entities = [{'start': sum(len(word) + 1 for word, _ in words_and_labels[:i]),
                             'end': sum(len(word) + 1 for word, _ in words_and_labels[:i + 1]), 'label': label} for
                            i, (word, label)
                            in enumerate(words_and_labels) if label != 'O']
                # print(entities)

                # Render the visualization without color for 'O' labels
                html = displacy.render(
                    {"text": " ".join(words), "ents": entities, "title": None},
                    style="ent",
                    manual=True,
                    options={"colors": {"MISC": "#806699",
                                        "ORG": "#ff6666",
                                        "LOC": "#66cc66",
                                        "PER": "#bf80ff",
                                        "O": None}}
                )
                # print(html)
                st.markdown(html, unsafe_allow_html=True)





        # Sử dụng widget st.html để hiển thị HTML

    # Hiển thị văn bản đã nhập
    # st.write("Văn bản đã nhập:", text)