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 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 = SentenceDetectorDLModel.pretrained("sentence_detector_dl", "xx").setInputCols(["document"]).setOutputCol("sentence")
    
    tokenizer_input = "sentence" if model != "bert_token_classifier_scandi_ner" else "document"
    tokenizer = Tokenizer().setInputCols([tokenizer_input]).setOutputCol("token")
    
    embeddings = {
        "norne_840B_300": 'glove_840B_300',
        "norne_6B_300": 'glove_6B_300',
        "norne_6B_100": 'glove_100d'
    }.get(model, 'glove_100d')

    embeddings_stage = WordEmbeddingsModel.pretrained(embeddings, lang='xx' if 'glove' in embeddings else None) \
        .setInputCols(['document', 'token']) \
        .setOutputCol('embeddings')

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

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

    stages = [document_assembler, sentence_detector, tokenizer, embeddings_stage, ner_model, ner_converter]
    return Pipeline(stages=stages)

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">Gjenkjenne personer, steder, organisasjoner og andre enheter</div>', unsafe_allow_html=True)

# Sidebar content
model = st.sidebar.selectbox(
    "Choose the pretrained model",
    ["bert_token_classifier_scandi_ner", "norne_840B_300", "norne_6B_300", "norne_6B_100"],
    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_NO.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 (født 28. oktober 1955) er en amerikansk forretningsmagnat, programvareutvikler, investor og filantrop. Han er best kjent som medgrunnlegger av Microsoft Corporation. I løpet av karrieren hos Microsoft hadde Gates stillingene som styreleder, administrerende direktør (CEO), president og sjefarkitekt for programvare, samtidig som han også var den største enkelteieren frem til mai 2014. Han er en av de mest kjente gründerne og pionerene innen mikrodatamaskinrevolusjonen på 1970- og 1980-tallet. Gates ble født og oppvokst i Seattle, Washington, og grunnla Microsoft sammen med barndomsvennen Paul Allen i 1975, i Albuquerque, New Mexico; det ble verdens største selskap for personlig datamaskinprogramvare. Gates ledet selskapet som styreleder og CEO fram til han trakk seg som CEO i januar 2000, men han forble styreleder og ble sjefarkitekt for programvare. På slutten av 1990-tallet ble Gates kritisert for sine forretningsmetoder, som ble ansett som anti-konkurransedyktige. Denne oppfatningen har blitt støttet av flere rettsavgjørelser. I juni 2006 annonserte Gates at han skulle gå over til en deltidsrolle i Microsoft og fulltidsarbeid ved Bill & Melinda Gates Foundation, den private veldedige stiftelsen han og kona Melinda Gates opprettet i 2000. Han overførte gradvis sine oppgaver til Ray Ozzie og Craig Mundie. Han trakk seg som styreleder i Microsoft i februar 2014 og overtok en ny stilling som teknologirådgiver for å støtte den nyutnevnte CEO-en Satya Nadella.",
    "Mona Lisa er et oljemaleri fra det 16. århundre skapt av Leonardo. Den er utstilt på Louvre i Paris.",
    "Da Sebastian Thrun begynte å jobbe med selvkjørende biler hos Google i 2007, tok få mennesker utenfor selskapet ham på alvor. 'Jeg kan fortelle deg at svært høytstående administrerende direktører for store amerikanske bilfirmaer ville håndhilse på meg og snu seg bort fordi det ikke var verdt å snakke med meg,' sa Thrun, nå medgrunnlegger og CEO for nettbasert høyere utdanningsoppstart Udacity, i et intervju med Recode tidligere denne uken.",
    "Facebook er en sosial nettverkstjeneste lansert som TheFacebook 4. februar 2004. Den ble grunnlagt av Mark Zuckerberg sammen med sine universitetskamerater og medstudenter ved Harvard University Eduardo Saverin, Andrew McCollum, Dustin Moskovitz og Chris Hughes. Medlemskapet på nettstedet var opprinnelig begrenset til Harvard-studenter, men ble utvidet til andre høyskoler i Boston-området, Ivy League og gradvis de fleste universiteter i USA og Canada.",
    "Historien om naturlig språkbehandling begynte generelt sett på 1950-tallet, selv om arbeid kan finnes fra tidligere perioder. I 1950 publiserte Alan Turing en artikkel med tittelen 'Computing Machinery and Intelligence' som foreslo det som nå kalles Turing-testen som kriterium for intelligens.",
    "Geoffrey Everest Hinton er en engelsk-kanadisk kognitiv psykolog og datavitenskapsmann, mest kjent for sitt arbeid med kunstige nevrale nettverk. Siden 2013 deler han tiden sin mellom Google og University of Toronto. I 2017 grunnla han og ble Chief Scientific Advisor for Vector Institute in Toronto.",
    "Da jeg fortalte John at jeg ville flytte til Alaska, advarte han meg om at jeg ville få problemer med å finne en Starbucks der.",
    "Steven Paul Jobs var en amerikansk forretningsmagnat, industriell designer, investor og medieeier. Han var styreleder, administrerende direktør (CEO) og medgrunnlegger av Apple Inc., styreleder og majoritetseier av Pixar, medlem av styret i The Walt Disney Company etter oppkjøpet av Pixar, og grunnlegger, styreleder og CEO for NeXT. Jobs er bredt anerkjent som en pioner innen den personlige datamaskinrevolusjonen på 1970- og 1980-tallet, sammen med Apple-medgrunnleggeren Steve Wozniak. Jobs ble født i San Francisco, California, og ble satt opp for adopsjon. Han ble oppdratt i San Francisco Bay Area. Han gikk på Reed College i 1972 før han droppet ut samme år, og reiste gjennom India i 1974 på jakt etter opplysning og studerte zen-buddhisme.",
    "Titanic er en amerikansk episk romantisk og katastrofefilm fra 1997 regissert,skrevet, co-produsert og co-redigert av James Cameron. Ved å inkludere både historiske og fiktive aspekter, er den basert på beretninger om RMS Titanics forlis, og har Leonardo DiCaprio og Kate Winslet i rollene som medlemmer av forskjellige sosiale klasser som forelsker seg om bord på skipet under dets ulykksalige jomfrureise.",
    "Bortsett fra å være kongen av nord, er John Snow en engelsk lege og en leder innen utviklingen av anestesi og medisinsk hygiene. Han anses for å være den første som brukte data for å kurere kolerautbruddet i 1834."
]

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)