import gradio as gr
import torch
import os

from model.process_text import process_text_pipeline
from pipeline import KeywordExtractorPipeline

DIR_PATH = os.path.dirname(os.path.realpath(__file__))


def extract_keyword(title, text, top_n, ngram_low_range, ngram_high_range, diversify_result):
    inp = {"text": text, "title": title}
    keyword_ls = kw_pipeline(inputs=inp, min_freq=1, ngram_n=(ngram_low_range, ngram_high_range),
                             top_n=top_n, diversify_result=diversify_result)
    result = ''
    for kw, score in keyword_ls:
        result += f'{kw}: {score}\n'
    highlight_dict = get_highlighted_text(text, keyword_ls)
    return result, highlight_dict


def get_highlighted_text(text, keyword_ls):
    text = process_text_pipeline(text)
    entities = []
    for kw, score in keyword_ls:

        kw = kw.replace('_', ' ')
        kw_dict = {"entity": "KEY"}
        start_idx = text.find(kw)

        while start_idx != -1:
            end_idx = start_idx + len(kw)
            kw_dict["start"] = start_idx
            kw_dict["end"] = end_idx
            entities.append(kw_dict)
            kw_dict = {"entity": "KEY"}
            start_idx = text.find(kw, end_idx)

    return {"text": text, "entities": entities}


if gr.NO_RELOAD:
    print("Loading PhoBERT model")
    phobert = torch.load(f'{DIR_PATH}/pretrained-models/phobert.pt')
    phobert.eval()

    print("Loading NER model")
    ner_model = torch.load(f'{DIR_PATH}/pretrained-models/ner-vietnamese-electra-base.pt')
    ner_model.eval()
    kw_pipeline = KeywordExtractorPipeline(phobert, ner_model)

if __name__ == "__main__":
    demo = gr.Interface(fn=extract_keyword,
                        inputs=[
                            gr.Text(
                                label="✍️ Title",
                                lines=1,
                                value="Truyền thuyết và hiện tại Thành Cổ Loa",
                                # value="Enter title here"
                            ),
                            gr.Textbox(
                                label="✍️ Text",
                                lines=5,
                                # value="Enter text here"
                                value="""Nhắc đến Cổ Loa, người ta nghĩ ngay đến truyền thuyết về An Dương Vương được thần Kim Quy bày cho cách xây thành, về chiếc lẫy nỏ thần làm từ móng chân rùa thần và mối tình bi thương Mỵ Châu – Trọng Thủy. Đằng sau những câu chuyện thiên về tâm linh ấy, thế hệ con cháu còn khám phá được những giá trị khảo cổ to lớn của Cổ Loa.
Khu di tích Cổ Loa cách trung – tâm Hà Nội 17km thuộc huyện Đông Anh, Hà Nội, có diện tích bảo tồn gần 500ha được coi là địa chỉ văn hóa đặc biệt của thủ đô và cả nước. Cổ Loa có hàng loạt di chỉ khảo cổ học đã được phát hiện, phản ánh quá trình phát triển liên tục của dân tộc ta từ sơ khai qua các thời kỳ đồ đồng, đồ đá và đồ sắt mà đỉnh cao là văn hóa Đông Sơn, vẫn được coi là nền văn minh sông Hồng thời kỳ tiền sử của dân tộc Việt Nam.
Cổ Loa từng là kinh đô của nhà nước Âu Lạc thời kỳ An Dương Vương (thế kỷ III TCN) và của nước Đại Việt thời Ngô Quyền (thế kỷ X) mà thành Cổ Loa là một di tích minh chứng còn lại cho đến ngày nay. Thành Cổ Loa được các nhà khảo cổ học đánh giá là “tòa thành cổ nhất, quy mô lớn vào bậc nhất, cấu trúc cũng thuộc loại độc đáo nhất trong lịch sử xây dựng thành lũy của người Việt cổ”.""",
                            ),
                            gr.Number(
                                label="Top N keywords",
                                info="Number of keywords retrieved",
                                value=5
                            ),
                            gr.Number(
                                label="Ngram low range",
                                value=1
                            ),
                            gr.Number(
                                label="Ngram high range",
                                value=3
                            ),
                            gr.Checkbox(
                                label="Diversify result"
                            )
                            ],
                        outputs=[gr.Textbox(label="Keywords Extracted"),
                                 gr.HighlightedText(label="Highlight",
                                                    color_map={"KEY":"blue"})],
                        title="🔑 KeyBERTVi",
                        description="Keyword Extraction for Vietnamese language",
#                         examples=[
#                             ["Truyền thuyết và hiện tại Thành Cổ Loa", """Nhắc đến Cổ Loa, người ta nghĩ ngay đến truyền thuyết về An Dương Vương được thần Kim Quy bày cho cách xây thành, về chiếc lẫy nỏ thần làm từ móng chân rùa thần và mối tình bi thương Mỵ Châu – Trọng Thủy. Đằng sau những câu chuyện thiên về tâm linh ấy, thế hệ con cháu còn khám phá được những giá trị khảo cổ to lớn của Cổ Loa.
# Khu di tích Cổ Loa cách trung – tâm Hà Nội 17km thuộc huyện Đông Anh, Hà Nội, có diện tích bảo tồn gần 500ha được coi là địa chỉ văn hóa đặc biệt của thủ đô và cả nước. Cổ Loa có hàng loạt di chỉ khảo cổ học đã được phát hiện, phản ánh quá trình phát triển liên tục của dân tộc ta từ sơ khai qua các thời kỳ đồ đồng, đồ đá và đồ sắt mà đỉnh cao là văn hóa Đông Sơn, vẫn được coi là nền văn minh sông Hồng thời kỳ tiền sử của dân tộc Việt Nam.
# Cổ Loa từng là kinh đô của nhà nước Âu Lạc thời kỳ An Dương Vương (thế kỷ III TCN) và của nước Đại Việt thời Ngô Quyền (thế kỷ X) mà thành Cổ Loa là một di tích minh chứng còn lại cho đến ngày nay. Thành Cổ Loa được các nhà khảo cổ học đánh giá là “tòa thành cổ nhất, quy mô lớn vào bậc nhất, cấu trúc cũng thuộc loại độc đáo nhất trong lịch sử xây dựng thành lũy của người Việt cổ”.""", 5, 1, 3, False],
#                             ["Nguy cơ ô nhiễm từ rác thực phẩm không phân loại", """Nếu như ở ngoại thành Hà Nội, mô hình phân loại và xử lý rác hữu cơ đã bắt đầu phát huy tác dụng thì trong nội thành rác thải từ thực phẩm, chủ yếu là thức ăn thừa bỏ đi, chiếm một lượng lớn trong thùng rác mỗi gia đình. Từ đó, tiềm ẩn nguy cơ gây ô nhiễm môi trường.
# Theo hướng dẫn kỹ thuật về phân loại chất thải rắn sinh hoạt" của Bộ TN&MT, rác thải từ thực phẩm là một trong số những loại bắt buộc người dân phải phân loại tại nguồn.
# TS. Vũ Thị Kim Tuyến – Hội Bảo vệ Thiên nhiên và Môi trường Việt Nam là một chuyên gia về lĩnh vực môi trường nên từ lâu đã hình thành cho mình và người thân thói quen phân loại rác tại nguồn. Trong đó, bà đặc biệt chú ý đến việc phân loại rác thải từ thực phẩm. Thế nhưng, tại nơi bà sống, mặc dù đã có quy định và các hộ dân cũng được tuyên truyền nhiều lần nhưng không phải ai cũng thực hiện phân loại như vậy, bởi nhiều lý do.
# Nếu như ở ngoại thành, mô hình phân loại và xử lý rác hữu cơ không chỉ giúp giảm đáng kể lượng rác thải ra môi trường trên địa bàn mà còn giúp người dân tiết kiệm được một khoản trong tưới tiêu, trồng trọt… thì ở nội thành, rác thực phẩm, chủ yếu là thức ăn thừa bỏ đi, chiếm một lượng lớn trong thùng rác mỗi gia đình. Thậm chí, ngay tại các khu chợ, rác thực phẩm cũng đang bị vứt bừa bãi, gây ô nhiễm môi trường, mất mỹ quan đô thị.""", 5, 1, 3, False]
#                         ]
                        )

    demo.launch(share=True)  # Share your demo with just 1 extra parameter 🚀