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

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

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

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

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

    pipeline = Pipeline(stages=[document_assembler, tokenizer, 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 Bengali text</div>', unsafe_allow_html=True)
st.markdown('<div class="section"><p>Recognize Persons, Locations, Organization, and Time entities in Bengali text using Deep Learning Model</p></div>', unsafe_allow_html=True)

# Sidebar content
model = st.sidebar.selectbox(
    "Choose the pretrained model",
    ["bengaliner_cc_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/NER_HINDI_ENGLISH.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 = [
    "উইলিয়াম হেনরি গেটস III (জন্ম অক্টোবর ২৮, ১৯৫৫) একজন আমেরিকান ব্যবসায়ী, সফটওয়্যার ডেভেলপার, বিনিয়োগকারী এবং দাতা। তিনি মাইক্রোসফট কর্পোরেশনের সহ-প্রতিষ্ঠাতা হিসেবে সর্বাধিক পরিচিত। মাইক্রোসফটে তাঁর কর্মজীবনে, গেটস চেয়ারম্যান, প্রধান নির্বাহী কর্মকর্তা (CEO), প্রেসিডেন্ট এবং প্রধান সফটওয়্যার স্থপতি হিসেবে দায়িত্ব পালন করেছেন, পাশাপাশি ২০১৪ সালের মে পর্যন্ত সর্বাধিক শেয়ারহোল্ডার ছিলেন। তিনি ১৯৭০ এবং ১৯৮০ এর দশকের মাইক্রোকম্পিউটার বিপ্লবের অন্যতম পরিচিত উদ্যোক্তা এবং অগ্রদূত। সিয়াটল, ওয়াশিংটনে জন্মগ্রহণ ও বড় হওয়া গেটস, ১৯৭৫ সালে শৈশবের বন্ধু পল অ্যালেনের সাথে আলবুকার্কি, নিউ মেক্সিকোতে মাইক্রোসফটের সহ-প্রতিষ্ঠা করেন; এটি বিশ্বের বৃহত্তম ব্যক্তিগত কম্পিউটার সফটওয়্যার কোম্পানিতে পরিণত হয়। গেটস ২০০০ সালের জানুয়ারি পর্যন্ত চেয়ারম্যান ও CEO হিসেবে কোম্পানিকে নেতৃত্ব দেন, কিন্তু তিনি চেয়ারম্যান হিসাবে রয়ে যান এবং প্রধান সফটওয়্যার স্থপতি হন। ১৯৯০-এর দশকের শেষের দিকে, গেটসের ব্যবসায়িক কৌশলগুলোর জন্য সমালোচিত হয়েছিলেন, যা প্রতিযোগিতাহীন বলে বিবেচিত হয়েছে। এই মতামতটি বিভিন্ন আদালতের রায় দ্বারা সমর্থিত হয়েছে। ২০০৬ সালের জুনে, গেটস ঘোষণা করেছিলেন যে তিনি মাইক্রোসফটে আংশিক সময়ের ভূমিকা এবং বিল ও মেলিন্ডা গেটস ফাউন্ডেশনে পূর্ণকালীন কাজ করবেন, যা তিনি এবং তাঁর স্ত্রী মেলিন্ডা গেটস ২০০০ সালে প্রতিষ্ঠা করেছিলেন। তিনি ধীরে ধীরে তার দায়িত্বগুলি রে ওজি এবং ক্রেইগ মন্ডিকে হস্তান্তর করেন। তিনি ২০১৪ সালের ফেব্রুয়ারিতে মাইক্রোসফটের চেয়ারম্যান পদ থেকে সরে যান এবং নতুন নিয়োগপ্রাপ্ত CEO সত্য নাদেলার সমর্থনে প্রযুক্তি উপদেষ্টা হিসাবে নতুন পদ গ্রহণ করেন।",
    "মোনা লিসা হল ১৬শ শতাব্দীর একটি তেলচিত্র যা লিওনার্দো দ্বারা নির্মিত। এটি প্যারিসের লুভরে সংরক্ষিত রয়েছে।",
    "যখন সেবাস্টিয়ান থ্রুন ২০০৭ সালে গুগলে স্বয়ংচালিত গাড়ির উপর কাজ শুরু করেছিলেন, তখন কোম্পানির বাইরে খুব কম লোকই তাকে গুরুত্ব সহকারে নিয়েছিলেন। 'আমি বলতে পারি যে প্রধান আমেরিকান গাড়ি কোম্পানির অনেক শীর্ষস্থানীয় CEO আমার সাথে হাত মিলাতেন এবং দূরে সরে যেতেন কারণ আমি কথা বলার যোগ্য ছিলাম না,' থ্রুন বলেছিলেন, বর্তমানে অনলাইন উচ্চ শিক্ষা স্টার্টআপ উদাসিটির সহ-প্রতিষ্ঠাতা এবং CEO, এই সপ্তাহের শুরুতে রিকোডের সাথে একটি সাক্ষাৎকারে।",
    "ফেসবুক একটি সামাজিক নেটওয়ার্কিং পরিষেবা যা ৪ ফেব্রুয়ারি, ২০০৪ তারিখে দ্যফেসবুক নামে চালু হয়। এটি হার্ভার্ড বিশ্ববিদ্যালয়ের ছাত্র মার্ক জুকারবার্গ তার কলেজ রুমমেট এবং সহপাঠী এডুয়ার্ডো স্যাভেরিন, অ্যান্ড্রু ম্যাককলাম, ডাস্টিন মস্কোভিটজ এবং ক্রিস হিউজের সাথে প্রতিষ্ঠা করেছিলেন। ওয়েবসাইটটির সদস্যপদ প্রাথমিকভাবে হার্ভার্ডের ছাত্রদের জন্য সীমাবদ্ধ ছিল, কিন্তু পরবর্তীতে বোস্টনের অন্যান্য কলেজ, আইভি লীগ এবং ধীরে ধীরে মার্কিন যুক্তরাষ্ট্র এবং কানাডার বেশিরভাগ বিশ্ববিদ্যালয়ে সম্প্রসারিত হয়।",
    "প্রাকৃতিক ভাষা প্রক্রিয়াকরণের ইতিহাস সাধারণত ১৯৫০ এর দশকে শুরু হয়েছিল, যদিও আগের সময়ের কাজ পাওয়া যায়। ১৯৫০ সালে, অ্যালান টুরিং একটি প্রবন্ধ প্রকাশ করেন 'কম্পিউটিং মেশিনারি অ্যান্ড ইন্টেলিজেন্স' শিরোনামে যা বুদ্ধিমত্তার মানদণ্ড হিসাবে এখন টুরিং পরীক্ষা নামে পরিচিত প্রস্তাব করেছিল।",
    "জিওফ্রি এভারেস্ট হিনটন একজন ইংরেজ-কানাডিয়ান কগনিটিভ সাইকোলজিস্ট এবং কম্পিউটার বিজ্ঞানী, কৃত্রিম নিউরাল নেটওয়ার্কে তার কাজের জন্য সর্বাধিক পরিচিত। ২০১৩ সাল থেকে তিনি গুগল এবং টরন্টো বিশ্ববিদ্যালয়ে কাজ করে তার সময় ভাগ করেন। ২০১৭ সালে, তিনি সহ-প্রতিষ্ঠা করেন এবং টরন্টোর ভেক্টর ইনস্টিটিউটের চিফ সায়েন্টিফিক অ্যাডভাইজার হন।",
    "যখন আমি জনকে বলেছিলাম যে আমি আলাস্কায় যেতে চাই, তিনি আমাকে সতর্ক করেছিলেন যে সেখানে স্টারবাকস খুঁজে পাওয়া আমার জন্য কঠিন হবে।",
    "স্টিভেন পল জবস একজন আমেরিকান ব্যবসায়ী, শিল্প ডিজাইনার, বিনিয়োগকারী এবং মিডিয়া মালিক ছিলেন। তিনি অ্যাপল ইনকর্পোরেটেডের চেয়ারম্যান, প্রধান নির্বাহী কর্মকর্তা (CEO) এবং সহ-প্রতিষ্ঠাতা ছিলেন, পিক্সারের চেয়ারম্যান এবং প্রধান শেয়ারহোল্ডার, এবং পিক্সার অধিগ্রহণের পর ওয়াল্ট ডিজনি কোম্পানির বোর্ডের সদস্য ছিলেন এবং নেক্সটের প্রতিষ্ঠাতা, চেয়ারম্যান এবং CEO ছিলেন। জবস ব্যাপকভাবে ব্যক্তিগত কম্পিউটার বিপ্লবের অগ্রদূত হিসেবে স্বীকৃত, অ্যাপলের সহ-প্রতিষ্ঠাতা স্টিভ ওজনিয়াকের সাথে। জবস সান ফ্রান্সিসকো, ক্যালিফোর্নিয়ায় জন্মগ্রহণ করেন এবং দত্তক হিসাবে লালিত হন। তিনি সান ফ্রান্সিসকো বে এলাকায় বেড়ে ওঠেন। ১৯৭২ সালে রিড কলেজে পড়াশোনা করেন এবং একই বছর পড়াশোনা ছেড়ে দেন, এবং ১৯৭৪ সালে ভারতের মধ্য দিয়ে ভ্রমণ করেন এবং আত্মশুদ্ধি এবং জেন বৌদ্ধধর্ম অধ্যয়ন করেন।",
    "টাইটানিক ১৯৯৭ সালের একটি আমেরিকান মহাকাব্যিক রোমান্স এবং বিপর্যয় চলচ্চিত্র যা জেমস ক্যামেরন পরিচালিত, লিখিত, সহ-প্রযোজিত এবং সহ-সম্পাদিত। এটি ঐতিহাসিক এবং কাল্পনিক দিকগুলোকে অন্তর্ভুক্ত করে এবং RMS টাইটানিকের ডুবে যাওয়ার ঘটনাগুলির উপর ভিত্তি করে নির্মিত, এবং এর প্রধান ভূমিকায় রয়েছেন লিওনার্দো ডিক্যাপ্রিও এবং কেট উইন্সলেট, যারা বিভিন্ন সামাজিক শ্রেণীর সদস্য যারা জাহাজের দুর্ভাগ্যজনক প্রথম যাত্রায় প্রেমে পড়ে।",
    "উত্তরের রাজা হওয়া ছাড়াও, জন স্নো একজন ইংরেজ চিকিৎসক এবং অ্যানেস্থেসিয়া এবং চিকিৎসা স্বাস্থ্যবিধির উন্নয়নের একজন নেতা। ১৮৩৪ সালে কলেরার প্রাদুর্ভাব নিরাময়ের জন্য ডেটা ব্যবহার করার জন্য তিনিই প্রথম হিসাবে বিবেচিত।"
]

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)