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

    if model == "ner_wikiner_glove_840B_300":
        embeddings = WordEmbeddingsModel.pretrained('glove_840B_300', lang='xx') \
            .setInputCols(['document', 'token']) \
            .setOutputCol('embeddings')
    elif model == "wikiner_6B_300":
        embeddings = WordEmbeddingsModel.pretrained('glove_6B_300', lang='xx') \
            .setInputCols(['document', 'token']) \
            .setOutputCol('embeddings')
    elif model == "wikiner_6B_100":
        embeddings = WordEmbeddingsModel.pretrained('glove_100d') \
            .setInputCols(['document', 'token']) \
            .setOutputCol('embeddings')

    lang = 'xx' if model == "ner_wikiner_glove_840B_300" else 'es'
    ner_model = NerDLModel.pretrained(model, lang) \
        .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">Reconocer personas, ubicaciones, organizaciones y entidades misceláneas</div>', unsafe_allow_html=True)

# Sidebar content
model = st.sidebar.selectbox(
    "Choose the pretrained model",
    ["ner_wikiner_glove_840B_300", "wikiner_6B_300", "wikiner_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_ES.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 = [
    """En el corazón de Madrid, la Puerta del Sol es uno de los lugares más emblemáticos de España. Salvador Dalí, uno de los pintores más célebres del siglo XX, nació en Figueres en 1904. A lo largo de su carrera, Dalí creó obras como "La persistencia de la memoria" y "Cristo de San Juan de la Cruz", que se exhiben en museos de Barcelona y Nueva York. El Museo del Prado, situado en el centro de Madrid, alberga obras maestras de artistas como Francisco de Goya y Diego Velázquez. El Rey Felipe VI, nacido Felipe Juan Pablo Alfonso de Todos los Santos de Borbón y de Grecia en Madrid en 1968, es el actual monarca de España. La Sagrada Familia, diseñada por Antoni Gaudí, es un símbolo de la arquitectura modernista y se encuentra en Barcelona. Pablo Picasso, nacido en Málaga en 1881, es famoso por sus pinturas y esculturas como "Guernica" y "Las señoritas de Aviñón", que se exhiben en el Museo Reina Sofía de Madrid. Sevilla es conocida por su catedral y el Alcázar. La Alhambra de Granada, un palacio y fortaleza de la época nazarí, es uno de los sitios más visitados del país. El Museo Guggenheim de Bilbao, diseñado por Frank Gehry, es un icono de la arquitectura contemporánea y alberga obras de artistas modernos. Valencia acoge cada año Las Fallas, un festival en el que se queman grandes figuras de cartón. Zaragoza es famosa por la Basílica del Pilar y el río Ebro. San Sebastián, con su festival de cine, atrae a celebridades de todo el mundo. En el ámbito literario, Miguel de Cervantes, autor de "Don Quijote de la Mancha", es una figura central de la literatura española. En la política, Pedro Sánchez, nacido en Madrid en 1972, es el actual presidente del Gobierno de España, mientras que el Partido Popular y el Partido Socialista Obrero Español son las principales fuerzas políticas del país. España también es conocida por su gastronomía, con platos típicos como la paella de Valencia, el gazpacho andaluz y el jamón ibérico de Salamanca. La selección española de fútbol, campeona del mundo en 2010, es una de las más exitosas en la historia del deporte.""",
    """William Henry Gates III (nacido el 28 de octubre de 1955) es un magnate de los negocios, desarrollador de software, inversor y filántropo estadounidense. Es mejor conocido como el cofundador de Microsoft Corporation. Durante su carrera en Microsoft, Gates ocupó los cargos de presidente, director ejecutivo (CEO), presidente y arquitecto de software en jefe, y también fue el mayor accionista individual hasta mayo de 2014. Es uno de los empresarios y pioneros más conocidos de revolución de la microcomputadora de los años setenta y ochenta. Nacido y criado en Seattle, Washington, Gates cofundó Microsoft con su amigo de la infancia Paul Allen en 1975, en Albuquerque, Nuevo México; se convirtió en la compañía de software de computadora personal más grande del mundo. Gates dirigió la compañía como presidente y CEO hasta que dejó el cargo de CEO en enero de 2000, pero siguió siendo presidente y se convirtió en el arquitecto jefe de software. A fines de la década de 1990, Gates había sido criticado por sus tácticas comerciales, que se han considerado anticompetitivas. Esta opinión ha sido confirmada por numerosas sentencias judiciales. En junio de 2006, Gates anunció que haría la transición a un puesto de medio tiempo en Microsoft y trabajaría a tiempo completo en la Fundación Bill y Melinda Gates, la fundación caritativa privada que él y su esposa, Melinda Gates, establecieron en 2000. [ 9] Poco a poco transfirió sus deberes a Ray Ozzie y Craig Mundie. Renunció como presidente de Microsoft en febrero de 2014 y asumió un nuevo cargo como asesor tecnológico para apoyar al recién nombrado CEO Satya Nadella.""",
    """La Mona Lisa es una pintura al óleo del siglo XVI creada por Leonardo. Se celebra en el Louvre de París.""",
    """Cuando Sebastian Thrun comenzó a trabajar en automóviles autónomos en Google en 2007, pocas personas fuera de la empresa lo tomaron en serio. "Puedo decirles que los CEOs muy importantes de las principales compañías de automóviles estadounidenses me darían la mano y se darían la vuelta porque no valía la pena hablar con ellos", dijo Thrun, ahora cofundador y CEO de la startup de educación superior en línea Udacity, en una entrevista. con Recode a principios de esta semana.""",
    """Facebook es un servicio de redes sociales lanzado como TheFacebook el 4 de febrero de 2004. Fue fundado por Mark Zuckerberg con sus compañeros de la universidad y sus compañeros de la Universidad de Harvard, Eduardo Saverin, Andrew McCollum, Dustin Moskovitz y Chris Hughes. La membresía del sitio web fue inicialmente limitada por los fundadores a los estudiantes de Harvard, pero se expandió a otras universidades en el área de Boston, la Ivy League, y gradualmente a la mayoría de las universidades en los Estados Unidos y Canadá.""",
    """La historia del procesamiento del lenguaje natural generalmente comenzó en la década de 1950, aunque se puede encontrar trabajo de períodos anteriores. En 1950, Alan Turing publicó un artículo titulado 'Maquinaria de computación e inteligencia' que proponía lo que ahora se llama la prueba de Turing como criterio de inteligencia""",
    """Geoffrey Everest Hinton es un psicólogo cognitivo y científico informático canadiense inglés, más conocido por su trabajo en redes neuronales artificiales. Desde 2013 divide su tiempo trabajando para Google y la Universidad de Toronto. En 2017, cofundó y se convirtió en el Asesor Científico Jefe del Vector Institute en Toronto.""",
    """Cuando le dije a John que quería mudarme a Alaska, me advirtió que tendría problemas para encontrar un Starbucks allí.""",
    """Steven Paul Jobs fue un magnate de los negocios, diseñador industrial, inversor y propietario de medios estadounidense. Fue presidente, director ejecutivo (CEO) y cofundador de Apple Inc., presidente y accionista mayoritario de Pixar, miembro de la junta directiva de The Walt Disney Company luego de la adquisición de Pixar, y fundador, presidente. y CEO de NeXT. Jobs es ampliamente reconocido como un pionero de la revolución de las computadoras personales de los años setenta y ochenta, junto con el cofundador de Apple, Steve Wozniak. Jobs nació en San Francisco, California, y fue puesto en adopción. Fue criado en el área de la bahía de San Francisco. Asistió a Reed College en 1972 antes de abandonar ese mismo año, y viajó por la India en 1974 buscando la iluminación y estudiando el budismo zen.""",
    """Titanic es una película de desastre y romance épico estadounidense de 1997 dirigida, escrita, coproducida y coeditada por James Cameron. Incorporando aspectos tanto históricos como ficticios, se basa en relatos del hundimiento del RMS Titanic, y está protagonizada por Leonardo DiCaprio y Kate Winslet como miembros de diferentes clases sociales que se enamoran a bordo del barco durante su desafortunado viaje inaugural.""",
    """Además de ser el rey del norte, John Snow es un médico inglés y un líder en el desarrollo de la anestesia y la higiene médica. Se lo considera el primero en usar datos para curar el brote de cólera en 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)