import streamlit as st
import sparknlp

from sparknlp.base import *
from sparknlp.annotator import *
from pyspark.ml import Pipeline

# 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 {
            background-color: #f9f9f9;
            padding: 10px;
            border-radius: 10px;
            margin-top: 10px;
        }
        .section p, .section ul {
            color: #666666;
        }
        .scroll {
            overflow-x: auto;
            border: 1px solid #e6e9ef;
            border-radius: 0.25rem;
            padding: 1rem;
            margin-bottom: 2.5rem;
            white-space: pre-wrap;
        }
    </style>
""", unsafe_allow_html=True)

@st.cache_resource
def init_spark():
    return sparknlp.start()

@st.cache_resource
def create_pipeline(model, task):
    documentAssembler = DocumentAssembler() \
        .setInputCol("text") \
        .setOutputCol("documents")

    t5 = T5Transformer.pretrained(model) \
        .setTask(task) \
        .setInputCols(["documents"]) \
        .setMaxOutputLength(200) \
        .setOutputCol("transfers")

    pipeline = Pipeline().setStages([documentAssembler, t5])
    return pipeline

def fit_data(pipeline, data):
    df = spark.createDataFrame([[data]]).toDF("text")
    result = pipeline.fit(df).transform(df)
    return result.select('transfers.result').collect()

# Sidebar setup
model = st.sidebar.selectbox(
    "Choose the Pretrained Model",
    ['t5_active_to_passive_styletransfer', 't5_passive_to_active_styletransfer'],
    help="Select the model you want to use for style transfer."
)

# Reference notebook link in sidebar
st.sidebar.markdown('Reference notebook:')
st.sidebar.markdown(
    """
    <a href="https://colab.research.google.com/github/JohnSnowLabs/spark-nlp-workshop/blob/master/tutorials/streamlit_notebooks/T5_LINGUISTIC.ipynb">
        <img src="https://colab.research.google.com/assets/colab-badge.svg" style="zoom: 1.3" alt="Open In Colab"/>
    </a>
    """, 
    unsafe_allow_html=True
)

examples = {
    "t5_active_to_passive_styletransfer": [
        "I am writing you a letter.",
        "Reporters write news reports.",
        "The company will hire new workers.",
        "Emma writes a letter.",
        "We did not grow rice.",
        "People will admire him.",
        "Someone has stolen my purse."
    ],
    "t5_passive_to_active_styletransfer": [
        "At dinner, six shrimp were eaten by Harry.",
        "The savannah is roamed by beautiful giraffes.",
        "The flat tire was changed by Sue.",
        "The students' questions are always answered by the teacher."
    ]
}

task_descriptions = {
    "t5_active_to_passive_styletransfer": "Transfer Active to Passive:",
    "t5_passive_to_active_styletransfer": "Transfer Passive to Active:"
}

# Set up the page layout
title = "Switch Between Active and Passive Voice"
sub_title = "Effortlessly Transform Sentences and Explore Different Writing Styles"

st.markdown(f'<div class="main-title">{title}</div>', unsafe_allow_html=True)
st.markdown(f'<div style="text-align: center; color: #666666;">{sub_title}</div>', unsafe_allow_html=True)

# Text selection and analysis
selected_text = st.selectbox("Select an example", examples[model])
custom_input = st.text_input("Try it with your own sentence!")

text_to_analyze = custom_input if custom_input else selected_text

st.write('Text to analyze:')
st.markdown(f'<div class="scroll">{text_to_analyze}</div>', unsafe_allow_html=True)

# Initialize Spark and create pipeline
spark = init_spark()
pipeline = create_pipeline(model, task_descriptions[model])
output = fit_data(pipeline, text_to_analyze)

# Display transformed sentence
st.write("Predicted Sentence:")
output_text = "".join(output[0][0])
st.markdown(f'<div class="scroll">{output_text.title()}</div>', unsafe_allow_html=True)