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")

    tokenizer = Tokenizer() \
        .setInputCols(["document"]) \
        .setOutputCol("token")

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

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

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

    nlpPipeline = Pipeline(stages=[documentAssembler, tokenizer, embeddings, ner, 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 Urdu text</div>', unsafe_allow_html=True)
st.markdown("""
<div class="section">
    <p>This demo utilizes embeddings-based NER model for Urdu texts, using the urduvec_140M_300d word embeddings</p>
</div>
""", unsafe_allow_html=True)

# Sidebar content
model = st.sidebar.selectbox(
    "Choose the pretrained model",
    ["uner_mk_140M_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_UR.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 = [
    """کھیل کے لیے سب سے مناسب تاریخ سال کا آخر ہے کیونکہ سنہ دوہزارباءیس کے آغاز میں سرمائی اولمپکس ہوتی ہیں۔ بلیٹر نے بین الاقوامی اولمپک کمیٹی کو یقین دلایا کہ ورلڈ کپ اور سرمائی اولمپکس کی تاریخوں میں کوئی ٹکراؤ نہیں ہو گا، جو کہ فروری دوہزارباءیس میں قزاقستان یا چین میں ہونی ہیں۔ فٹبال کلب کی تاریخ میں سنیچر 25اکتوبر کو ریئل میڈرڈ اور بارسلونا کے درمیان اہم ترین میچ کھیلا جائے گا۔ اس میچ میں فٹبال کی دنیا کے دو بہترین کھلاڑی حصہ لیں گے۔ اس میچ کی ایک خاص بات اور بھی ہے کہ اس میں فٹبال کے عالمی کپ میں سب سے زیادہ گول کرنے والے کھلاڑی کے علاوہ تاریخ کے پانچ میں سے چار مہنگے ترین کھلاڑی بھی حصہ لیں گے۔""",
    """اس میچ کی سب سے زیادہ سحر انگيزی یوروگوائے کے مشہور کھلاڑی لوئس سواریز کی شمولیت ہے۔ خیال رہے کہ لوئس سواریز نے برازیل میں کھیلے جانے والے فٹبال کے عالمی کپ کے ایک میچ کے دوران اطالوی کھلاڑی جیورجیو چیلینی کو کاٹ کھایا تھا جس کے بعد فیفا نے ان پر نو بین الاقوامی میچوں کی پابندی عائد کی تھی۔ اس میچ میں ارجنٹائن سے تعلق رکھنے والے فٹبال کے شہرۂ آفاق کھلاڑی لیونل میسی اور برازیل کے سٹار کھلاڑی نیمار بھی حصہ لیں گے۔""",
    """سائنسدانوں نے صرف بیس ڈالر کا مواد استعمال کرتے ہوئے صرف بارہ گھنٹے میں ایبولا کا پروٹو ٹائپ بنا کر اپنی ٹیسٹ کی تکنیک کو ثابت کر دیا۔ اس ٹیسٹ میں حیاتیاتی اجزا استعمال ہوتے ہیں جس میں آر این اے کے جنیاتی اجزا بھی شامل ہوتے ہیں۔ سائنسدانوں کا کہنا ہے کہ ان اجزا کو ٹھنڈا اور خشک کرکے عام کاغذ پر بھی سٹور کیا جا سکتا ہے۔ تحقیقی ٹیم کے سربراہ جیم کولنز جو بوسٹن اور ہارورڈ دونوں یونیورسٹیوں میں تعینات ہیں، کہتے ہیں کہ ان حیاتیاتی اجزا کو صرف پانی ملانے سے متحرک کیا جا سکتا ہے۔ انھوں نے بی بی سی کو بتایا کہ ہمیں اس بات پر حیرت ہوئی کہ ٹھنڈا اور خشک کرنے کے بعد ان اجزا نے کتنے اچھے طریقے سے کام کیا۔""",
    """ان کے بائیں ہاتھ میں ہیئر لائن فریکچر ہے اور وہ جمعے کو میدان میں فیلڈنگ کے لیے بھی نہیں آئے۔ تین میچوں کی اس ٹیسٹ سیریز میں پاکستان کو ایک صفر کی برتری حاصل ہے اور وہ سری لنکا کی سرزمین پر نو سال میں پہلی ٹیسٹ سیریز جیتنے کے لیے کوشاں ہے۔ پاکستانی ٹیم نے آخری بار سری لنکا میں ٹیسٹ سیریز سنہ 2006 میں انضمام الحق کی قیادت میں جیتی تھی جس کے بعد لگاتار تین مرتبہ اسے ٹیسٹ سیریز میں شکست ہو چکی ہے۔""",
    """انھوں نے کہا کہ گذشتہ ایشیز سیریز میں پانچ صفر کی جیت اور ورلڈ کپ کا فاتح بننا کلارک کے کریئر کے دو اہم واقعات ہیں۔ اس کے علاوہ بھی انھوں نے اپنے ملک کے لیے بحیثیت بیٹسمین اور کپتان غیرمعمولی کارکردگی کا مظاہرہ کیا۔ڈربن میں کھیلے جانے والے تیسرے ایک روزہ میچ میں جنوبی افریقہ نے نیوزی لینڈ کو 62 رنز سے شکست دے کر تین میچوں پر مشتمل سیریز دو ایک سے جیت لی۔"""
]

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)