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_distilbert_sentiment", "vi") \
    .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 Vietnamese Sentiment Detection with Spark NLP</div>', unsafe_allow_html=True)

# Sidebar content
model = st.sidebar.selectbox(
    "Choose the pretrained model",
    ["classifierdl_distilbert_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/SENTIMENT_EN.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 = [
    "Chủ shop nch cực nhiệt tình và dễ thương lắm luôn á, lại còn tự tay viết thiệp và tặng kèm quà nữa, hihi.",
    "Cho 3 sao vì đã nói trước là giúp mình cắt dây đeo là 6 6 rồi nhưng không cắt.",
    "Dày ko y hình vì nếu y hình thì chữ nike phải may bằng kim tuyến mình check mã thì ko có mặt hàng này dù rất buồn nhưng cx đã nhận hàng rồi và ko biết phải làm sao😑😑.",
    "Mặc dù có chút trục trặc về việc giao hàng nhưng chủ shop gọi điện nói chuyện rất lịch sự nên cũng cho qua được.",
    "Chất lượng sản phẩm rất kém Shop làm ăn ko uy tín gửi hàng ko đúng hình.",
    "Chất lượng sản phẩm tốt Đóng gói sản phẩm chắc chắn Thời gian giao hàng nhanh chỉ tiếc dây xanh gắn dây kéo bị ra màu nên e phải cắt bỏ nhưng sản phẩn khá ổn hợp với giá tiền."
]

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)