import streamlit as st
import sparknlp
import os
import pandas as pd

from sparknlp.base import *
from sparknlp.annotator import *
from pyspark.ml import Pipeline
from sparknlp.pretrained import PretrainedPipeline
from annotated_text import annotated_text

# Page configuration
st.set_page_config(
    layout="wide", 
    initial_sidebar_state="auto"
)

# CSS for styling
st.markdown("""
    <style>
        .main-title {
            font-size: 36px;
            color: #4A90E2;
            font-weight: bold;
            text-align: center;
        }
        .section {
            background-color: #f9f9f9;
            padding: 10px;
            border-radius: 10px;
            margin-top: 10px;
        }
        .section p, .section ul {
            color: #666666;
        }
    </style>
""", unsafe_allow_html=True)

@st.cache_resource
def init_spark():
    return sparknlp.start()

@st.cache_resource
def create_pipeline(model):
    documentAssembler = DocumentAssembler()\
    .setInputCol("text")\
    .setOutputCol("document")

    sentenceDetector = SentenceDetectorDLModel.pretrained("sentence_detector_dl", "xx")\
    .setInputCols(["document"])\
    .setOutputCol("sentence")

    tokenizer = WordSegmenterModel.pretrained("wordseg_large", "zh") \
        .setInputCols(["sentence"]) \
        .setOutputCol("token")

    tokenClassifier = XlmRoBertaForTokenClassification.pretrained("xlm_roberta_large_token_classifier_hrl", "xx")\
    .setInputCols(["sentence",'token'])\
    .setOutputCol("ner")

    ner_converter = NerConverter()\
    .setInputCols(["sentence", "token", "ner"])\
    .setOutputCol("ner_chunk")

    nlpPipeline = Pipeline(stages=[documentAssembler, sentenceDetector, tokenizer, tokenClassifier, ner_converter])
    return nlpPipeline

def fit_data(pipeline, data):
  empty_df = spark.createDataFrame([['']]).toDF('text')
  pipeline_model = pipeline.fit(empty_df)
  model = LightPipeline(pipeline_model)
  result = model.fullAnnotate(data)
  return result

def annotate(data):
    document, chunks, labels = data["Document"], data["NER Chunk"], data["NER Label"]
    annotated_words = []
    for chunk, label in zip(chunks, labels):
        parts = document.split(chunk, 1)
        if parts[0]:
            annotated_words.append(parts[0])
        annotated_words.append((chunk, label))
        document = parts[1]
    if document:
        annotated_words.append(document)
    annotated_text(*annotated_words)

# Set up the page layout
st.markdown('<div class="main-title">Recognize entities in Chinese text</div>', unsafe_allow_html=True)

# Sidebar content
model = st.sidebar.selectbox(
    "Choose the pretrained model",
    ["xlm_roberta_large_token_classifier_hrl"],
    help="For more info about the models visit: https://sparknlp.org/models"
)

# Reference notebook link in sidebar
link = """
<a href="https://colab.research.google.com/github/JohnSnowLabs/spark-nlp-workshop/blob/master/tutorials/streamlit_notebooks/NER.ipynb">
    <img src="https://colab.research.google.com/assets/colab-badge.svg" style="zoom: 1.3" alt="Open In Colab"/>
</a>
"""
st.sidebar.markdown('Reference notebook:')
st.sidebar.markdown(link, unsafe_allow_html=True)

# Load examples
examples = [
    "当 前 ， 在 中共十五大 精 神 的 指 引 下 ， 在 以 江泽民 同 志 为 核 心 的 中共中央 领 导 下 ， 全 党 和 全 国 各 族 人 民 正 高 举 邓小平 理 论 伟 大 旗 帜 ， 同 心 同 德 ， 团 结 奋 斗 ， 沿 着 建 设 有 中国 特 色 的 社 会 主 义 道 路 阔 步 前 进 。",
    "中共中央 致 中国致公党十一大 的 贺 词 各 位 代 表 、 各 位 同 志 ： 在 中国致公党第十一次全国代表大会 隆 重 召 开 之 际 ， 中国共产党中央委员会 谨 向 大 会 表 示 热 烈 的 祝 贺 ， 向 致公党 的 同 志 们 致 以 亲 切 的 问 候 ！",
    "数 百 名 华 人 、 华 侨 、 留 学 人 员 、 我 国 驻 纽约 总 领 事 馆 代 表 在 机 场 挥 舞 中 美 两 国 国 旗 ， 热 烈 欢 迎 江 主 席 访 问 波士顿 。",
    "到 机 场 迎 接 江 主 席 的 美 方 人 员 有 马萨诸塞州 州 长 和 波士顿 市 长 等 。",
    "又 讯 中国 国 家 主 席 江泽民 １ 日 上 午 应 邀 在 美国 著 名 学 府 哈佛大学 发 表 重 要 演 讲 。",
    "江 主 席 来 到 哈佛大学 时 ， 受 到 哈佛大学 校 长 陆登庭 及 哈佛 各 学 院 院 长 的 热 烈 欢 迎 。",
    "本 报 纽约 １ ０ 月 ３ １ 日 电 记 者 陈特安 、 周德武 报 道 ： 今 天 晚 上 ， 美中贸易全国委员会 和 美国中国商会 在 纽约 举 行 盛 大 宴 会 欢 迎 江泽民 主 席 。",
    "哈佛大学 校 长 陆登庭 对 江 主 席 访 问 哈佛 并 发 表 演 讲 表 示 欢 迎 。",
    "美中贸易全国委员会 主 席 费希尔 和 美国中国商会 会 长 沈被章 先 后 致 词 。"
]

selected_text = st.selectbox("Select an example", examples)
custom_input = st.text_input("Try it with your own Sentence!")

text_to_analyze = custom_input if custom_input else selected_text

st.subheader('Full example text')
HTML_WRAPPER = """<div class="scroll entities" style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; margin-bottom: 2.5rem; white-space:pre-wrap">{}</div>"""
st.markdown(HTML_WRAPPER.format(text_to_analyze), unsafe_allow_html=True)

# Initialize Spark and create pipeline
spark = init_spark()
pipeline = create_pipeline(model)
output = fit_data(pipeline, text_to_analyze)

# Display matched sentence
st.subheader("Processed output:")

results = {
    'Document': output[0]['document'][0].result,
    'NER Chunk': [n.result for n in output[0]['ner_chunk']],
    "NER Label": [n.metadata['entity'] for n in output[0]['ner_chunk']]
}

annotate(results)

with st.expander("View DataFrame"):
    df = pd.DataFrame({'NER Chunk': results['NER Chunk'], 'NER Label': results['NER Label']})
    df.index += 1
    st.dataframe(df)