import streamlit as st
import numpy as np
import tensorflow as tf
from transformers import RobertaTokenizerFast, TFRobertaModel

MAX_LEN = 96
MODEL_PATH = "src/roberta_fold_1.h5" 

#MODEL ARCHITECTURE
def build_model():
    ids = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32, name="input_ids")
    att = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32, name="attention_mask")

    roberta = TFRobertaModel.from_pretrained('roberta-base', from_pt=True)
    x = roberta(ids, attention_mask=att)[0]
    x1 = tf.keras.layers.Dropout(0.1)(x) 
    start_logits = tf.keras.layers.Dense(1)(x1)
    start_logits = tf.keras.layers.Flatten()(start_logits)
    end_logits = tf.keras.layers.Dense(1)(x1)
    end_logits = tf.keras.layers.Flatten()(end_logits)
    
    model = tf.keras.Model(inputs=[ids, att], outputs=[start_logits, end_logits])
    return model

#MODEL
@st.cache_resource
def load_model_and_tokenizer():
    tokenizer = RobertaTokenizerFast.from_pretrained('roberta-base')
    model = build_model()
    model.load_weights(MODEL_PATH)
    return model, tokenizer

with st.spinner('Loading model and tokenizer... Please wait.'):
    model, tokenizer = load_model_and_tokenizer()

#PREDICTION 
def predict_sentiment(tweet, sentiment):
    # Preprocessing
    tweet = " " + " ".join(str(tweet).split())
    # Tokenization
    enc = tokenizer.encode_plus(
        sentiment, 
        tweet,
        add_special_tokens=True,
        max_length=MAX_LEN,
        padding='max_length',
        truncation=True,
        return_attention_mask=True)
    
    input_ids = np.array([enc['input_ids']])
    attention_mask = np.array([enc['attention_mask']])
    # Inference
    preds = model.predict([input_ids, attention_mask])
    start_logits = preds[0]
    end_logits = preds[1]
    
    idx_start = np.argmax(start_logits)
    idx_end = np.argmax(end_logits)
    # Decoding
    if sentiment == "neutral" or len(tweet.split()) < 2:
        return tweet.strip()
    # Validation check
    if idx_start > idx_end:
        return tweet.strip()
    else:
        text_tokens = tokenizer.decode(enc['input_ids'][idx_start:idx_end+1])
        return text_tokens.strip()

#INTERFACE
st.title("🐦 Tweet Sentiment Extraction")
st.markdown("""
This AI model analyzes a tweet and extracts the **specific phrase or word** that supports the selected sentiment.
*Built with TensorFlow & RoBERTa*
""")

#Inputs
col1, col2 = st.columns([1, 2])
with col1:
    sentiment = st.selectbox(
        "Select Sentiment:", 
        ["positive", "negative", "neutral"])
with col2:
    text = st.text_area(
        "Enter Tweet:", 
        value="This movie was really fantastic and I loved it!",
        height=100)
#Prediction
if st.button("Extract Phrase", type="primary"):
    if text:
        result = predict_sentiment(text, sentiment)
        st.write("### Result:")
        # Highlight Logic for Visualization
        if result in text and len(result) > 2:
            highlighted_text = text.replace(result, f"<span style='background-color: #ffd700; color: black; padding: 2px; font-weight: bold; border-radius: 3px;'>{result}</span>")
            st.markdown(f"**Context:** {highlighted_text}", unsafe_allow_html=True)
        else:
            st.markdown(f"**Context:** {text}")  
        st.success(f"Extracted Text: **{result}**")
    else:
        st.warning("Please enter a tweet to analyze.")