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", 
    page_title="Spark NLP Demos App", 
    initial_sidebar_state="auto"
)

# CSS for styling
st.markdown("""
    <style>
        .main-title {
            font-size: 36px;
            color: #4A90E2;
            font-weight: bold;
            text-align: center;
        }
        .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('glove_840B_300', lang='xx') \
        .setInputCols(["document", "token"]) \
        .setOutputCol("embeddings")

    ner_model = NerDLModel.pretrained(model, 'xx') \
        .setInputCols(["document", "token", "embeddings"]) \
        .setOutputCol("ner")

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

    pipeline = Pipeline(stages=[
        document_assembler,
        tokenizer,
        embeddings,
        ner_model,
        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">Erkennen Sie Personen, Standorte, Organisationen und verschiedene Entitäten</div>', unsafe_allow_html=True)

# Sidebar content
model = st.sidebar.selectbox(
    "Choose the pretrained model",
    ["ner_wikiner_glove_840B_300"],
    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_DE.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 = [
    """William Henry Gates III (* 28. Oktober 1955 in London) ist ein US-amerikanischer Geschäftsmann, Softwareentwickler, Investor und Philanthrop. Er ist bekannt als Mitbegründer der Microsoft Corporation. Während seiner Karriere bei Microsoft war Gates Vorsitzender, Chief Executive Officer (CEO), Präsident und Chief Software Architect und bis Mai 2014 der größte Einzelaktionär. Er ist einer der bekanntesten Unternehmer und Pioniere der Mikrocomputer-Revolution der 1970er und 1980er Jahre. Gates wurde in Seattle, Washington, geboren und wuchs dort auf. 1975 gründete er Microsoft zusammen mit seinem Freund aus Kindertagen, Paul Allen, in Albuquerque, New Mexico. Es entwickelte sich zum weltweit größten Unternehmen für Personal-Computer-Software. Gates leitete das Unternehmen als Chairman und CEO, bis er im Januar 2000 als CEO zurücktrat. Er blieb jedoch Chairman und wurde Chief Software Architect. In den späten neunziger Jahren wurde Gates für seine Geschäftstaktiken kritisiert, die als wettbewerbswidrig angesehen wurden. Diese Meinung wurde durch zahlreiche Gerichtsurteile bestätigt. Im Juni 2006 gab Gates bekannt, dass er eine Teilzeitstelle bei Microsoft und eine Vollzeitstelle bei der Bill & Melinda Gates Foundation, der privaten gemeinnützigen Stiftung, die er und seine Frau Melinda Gates im Jahr 2000 gegründet haben, übernehmen wird. [ 9] Er übertrug seine Aufgaben nach und nach auf Ray Ozzie und Craig Mundie. Im Februar 2014 trat er als Vorsitzender von Microsoft zurück und übernahm eine neue Position als Technologieberater, um den neu ernannten CEO Satya Nadella zu unterstützen.""",
    """Die Mona Lisa ist ein Ölgemälde aus dem 16. Jahrhundert, das von Leonardo geschaffen wurde. Es findet im Louvre in Paris statt.""",
    """Als Sebastian Thrun 2007 bei Google anfing, an selbstfahrenden Autos zu arbeiten, nahmen ihn nur wenige Leute außerhalb des Unternehmens ernst. "Ich kann Ihnen sagen, dass sehr hochrangige CEOs großer amerikanischer Automobilunternehmen mir die Hand schütteln und sich abwenden würden, weil ich es nicht wert war, mit ihnen zu sprechen", sagte Thrun, jetzt Mitbegründer und CEO des Online-Hochschul-Startups Udacity, in einem Interview mit Recode Anfang dieser Woche.""",
    """Facebook ist ein sozialer Netzwerkdienst, der am 4. Februar 2004 als TheFacebook gestartet wurde. Er wurde von Mark Zuckerberg mit seinen Mitbewohnern und Kommilitonen der Harvard University, Eduardo Saverin, Andrew McCollum, Dustin Moskovitz und Chris Hughes, gegründet. Die Mitgliedschaft der Website wurde ursprünglich von den Gründern auf Harvard-Studenten beschränkt, aber auf andere Colleges in der Region Boston, die Ivy League und nach und nach auf die meisten Universitäten in den USA und Kanada ausgeweitet.""",
    """Die Geschichte der Verarbeitung natürlicher Sprache begann im Allgemeinen in den 1950er Jahren, obwohl Arbeiten aus früheren Perioden gefunden werden können. 1950 veröffentlichte Alan Turing einen Artikel mit dem Titel "Computing Machinery and Intelligence", in dem der heutige Turing-Test als Kriterium für die Intelligenz vorgeschlagen wurde""",
    """Geoffrey Everest Hinton ist ein englisch-kanadischer kognitiver Psychologe und Informatiker, der vor allem für seine Arbeit an künstlichen neuronalen Netzen bekannt ist. Seit 2013 arbeitet er für Google und die University of Toronto. 2017 war er Mitbegründer und wissenschaftlicher Leiter des Vector Institute in Toronto.""",
    """Als ich John sagte, dass ich nach Alaska ziehen wollte, warnte er mich, dass ich dort Probleme haben würde, einen Starbucks zu finden.""",
    """Steven Paul Jobs war ein amerikanischer Geschäftsmagnat, Industriedesigner, Investor und Medieninhaber. Er war Vorsitzender, Chief Executive Officer (CEO) und Mitbegründer von Apple Inc., Vorsitzender und Mehrheitsaktionär von Pixar, einem Mitglied des Board of Directors der Walt Disney Company nach der Übernahme von Pixar, und Gründer, Vorsitzender und CEO von NeXT. Jobs ist weithin als Pionier der PC-Revolution der 1970er und 1980er Jahre anerkannt, zusammen mit Apple-Mitbegründer Steve Wozniak. Jobs wurde in San Francisco, Kalifornien, geboren und zur Adoption freigegeben. Er wuchs in der San Francisco Bay Area auf. Er besuchte das Reed College 1972, bevor er es im selben Jahr abbrach, und reiste 1974 durch Indien, um Erleuchtung zu erlangen und den Zen-Buddhismus zu studieren.""",
    """Titanic ist ein amerikanischer epischer Romantik- und Katastrophenfilm aus dem Jahr 1997, der von James Cameron inszeniert, geschrieben, co-produziert und mitherausgegeben wurde. Es enthält sowohl historische als auch fiktive Aspekte und basiert auf Berichten über den Untergang der RMS Titanic. Die Stars Leonardo DiCaprio und Kate Winslet sind Mitglieder verschiedener sozialer Schichten, die sich während ihrer unglücklichen Jungfernfahrt an Bord des Schiffes verlieben.""",
    """John Snow ist nicht nur der König des Nordens, sondern auch ein englischer Arzt und führend in der Entwicklung von Anästhesie und medizinischer Hygiene. Er gilt als der erste, der Daten zur Heilung des Cholera-Ausbruchs im Jahr 1834 verwendet."""
]

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)