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

# 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 = DocumentAssembler()\
        .setInputCol("text")\
        .setOutputCol("document")

    embeddings = BertSentenceEmbeddings\
        .pretrained('labse', 'xx') \
        .setInputCols(["document"])\
        .setOutputCol("sentence_embeddings")

    sentimentClassifier = ClassifierDLModel.pretrained("classifierdl_bert_sentiment", "es") \
    .setInputCols(["sentence_embeddings"]) \
    .setOutputCol("class_")  
        
    nlpPipeline = Pipeline(
        stages=[
            document, 
            embeddings,
            sentimentClassifier])

    return nlpPipeline

def fit_data(pipeline, data):
    empty_df = spark.createDataFrame([['']]).toDF('text')
    pipeline_model = pipeline.fit(empty_df)
    model = LightPipeline(pipeline_model)
    results = model.fullAnnotate(data)[0]

    return results['class_'][0].result

# Set up the page layout
st.markdown('<div class="main-title">State-of-the-Art Spanish Sentiment Detection with Spark NLP</div>', unsafe_allow_html=True)

# Sidebar content
model = st.sidebar.selectbox(
    "Choose the pretrained model",
    ["classifierdl_bert_sentiment"],
    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/BertForSequenceClassification.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 = [
    "Hr.formand, kære kommissær og kære kolleger, jeg vil starte med at sige Hr.Rapkay tak for en god betænkning og et godt samarbejde.",
    "Protester over hele landet ledet af utilfredse civilsamfund på grund af den danske regerings COVID-19 lockdown-politik er kommet ud af kontrol.",
    "Hvidbogen repræsenterer tre til fire måneders intensivt arbejde, siden den nye kommission blev udpeget i september.",
    "Det er ikke et nemt emne, og jeg mener derfor, at den indsats, som fru Lienemanns har ydet, fortjener stor respekt.",
    "Det er desværre en sørgelig erkendelse af de store vanskeligheder, som Wales oplever.",
    "Vi har indarbejdet resultaterne af de omfattende høringer, der har fundet sted i løbet af det sidste par år siden offentliggørelsen af kommissionens grønbog om fødevarelovgivning.",
    "Fru formand, jeg vil gerne først give Dem en kompliment for den kendsgerning, at De har holdt Deres ord, og at antallet af tv-kanaler på vores kontorer faktisk er udvidet enormt nu i denne første mødeperiode i det nye år.",
    "Det er fuldstændig skandaløst, at vi fastsætter lovgivningsbestemmelser og så ikke overholder dem selv.",
    "Min gruppe har udarbejdet omfattende ændringsforslag til begge betænkninger til forhandling i dag."
        ]

selected_text = st.selectbox("Select a sample", examples)
custom_input = st.text_input("Try it for yourself!")

if custom_input:
    selected_text = custom_input
elif selected_text:
    selected_text = selected_text

st.subheader('Selected Text')
st.write(selected_text)

# Initialize Spark and create pipeline
spark = init_spark()
pipeline = create_pipeline(model)
output = fit_data(pipeline, selected_text)

# Display output sentence
if output.lower() in ['pos', 'positive']:
  st.markdown("""<h3>This seems like a <span style="color: green">{}</span> text. <span style="font-size:35px;">&#128515;</span></h3>""".format('positive'), unsafe_allow_html=True)
elif output.lower() in ['neg', 'negative']:
  st.markdown("""<h3>This seems like a <span style="color: red">{}</span> text. <span style="font-size:35px;">&#128544;</span?</h3>""".format('negative'), unsafe_allow_html=True)