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):
    document_assembler = DocumentAssembler() \
        .setInputCol("text") \
        .setOutputCol("document")

    sentence_detector = SentenceDetector() \
        .setInputCols(["document"]) \
        .setOutputCol("sentence")

    word_segmenter = WordSegmenterModel.pretrained("wordseg_kaist_ud", "ko") \
        .setInputCols(["sentence"]) \
        .setOutputCol("token")

    embeddings = WordEmbeddingsModel.pretrained("glove_840B_300", "xx") \
        .setInputCols(["document", "token"]) \
        .setOutputCol("embeddings")

    ner = NerDLModel.pretrained("ner_kmou_glove_840B_300d", "ko") \
        .setInputCols(["document", "token", "embeddings"]) \
        .setOutputCol("ner")

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

    pipeline = Pipeline(stages=[document_assembler, sentence_detector, word_segmenter, embeddings, ner, ner_converter])
    return pipeline

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 Urdu text</div>', unsafe_allow_html=True)
st.markdown("""
<div class="section">
    <p>This model uses the pre-trained <code>glove_840B_300</code> embeddings model from WordEmbeddings annotator as an input</p>
</div>
""", unsafe_allow_html=True)

# Sidebar content
model = st.sidebar.selectbox(
    "Choose the pretrained model",
    ["ner_kmou_glove_840B_300d"],
    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/public/NER_KO.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 = [
    """ARD , ZDF 등 공영 TV 와 바이에른주 방송 , 북부 독일 방송 등 은 이 날 한국 의 총선 소식 과 관련 , 여당 의 과반수 의석 확보 와 신당 의 득표 율 이 이번 선거 의 최대 관심사 이 라고 보도 하 ㄴ 데 잇 어 저녁 시간 부터 는 수 차례 에 걸치 어 개표 상황 과 정당 별 의석 전망 을 속보 로 전하 았 다 .""",
    """두 나라 관계 는 중국 의 인권 문제 와 핵확산 방지 문제 , 통상 문제 및 최근 의 F 16 전투기 대 대만 판매 등 을 놓 고 이미 위험선 상 에 오 아 있 는데 클린턴 행정부 의 등장 으로 양국 관계 가 더욱 경색 되 ㄹ 것 을 걱정 하 는 분위기 .""",
    """서울대 건축공학 과 를 졸업 하 ㄴ 이 씨 는 한국건축가협회""",
    """나 는 다시 순자 를 양동 에서 빼내 기 위하 아서 창신 팔동""",
    """헤라신전 서 채화 되 ㄴ 지 보름 , 지구 의 반바퀴 를 돌 아 제주공항 에 첫발 을 내디디 ㄴ 이래 로 열이틀""",
    """다음 은 홍콩 의 권위지 명보 와 일본 도쿄 ( 동경 ) 신문 이 24일""",
    """최 영사 가 우리 외교관 이 며 그 신변보호 책임 이 주재국 이 ㄴ 러시아 에 있 다는 점 에서 러시아 는 이 같 은 우리 정부 요구 에 응하 아야 하 ㄹ 의무 가 있 다 .""",
    """판 에 박 은 듯 하 ㄴ 깨끗 하 ㄴ 글씨 로 , 처음 단군 님 이 니 신라 , 백제 , 고구려 이 니 띄엄띄엄 어른 들 한테 서 귀결 로 들어오 던 얘기 들 이 참말 로 씌 어 있 었 다 ."""
]

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)