import gradio as gr
import joblib
import pandas as pd
import numpy as np
import re
import nltk
import pickle
import tensorflow as tf
from nltk.corpus import stopwords
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.sequence import pad_sequences

# --- Setup & Preprocessing ---
nltk.download('stopwords')
stop_words = set(stopwords.words('english'))

def preprocess(text):
    text = str(text).lower()
    text = re.sub(r'[^a-zA-Z\s]', '', text)
    tokens = text.split()
    return " ".join([w for w in tokens if w not in stop_words])

# --- Load Artifacts ---
print("Loading models...")

# 1. Load Scikit-Learn Models & Vectorizer
try:
    lr_model = joblib.load('model_lr.joblib')
    nb_model = joblib.load('model_nb.joblib')
    tfidf = joblib.load('tfidf_vectorizer.joblib')
except Exception as e:
    print(f"Error loading Sklearn models: {e}")

# 2. Load Keras Models & Tokenizer
try:
    ffnn_model = load_model('model_ffnn.h5')
    embed_model = load_model('model_embed.h5')
    with open('tokenizer.pickle', 'rb') as handle:
        tokenizer = pickle.load(handle)
except Exception as e:
    print(f"Error loading Keras models/tokenizer: {e}")

# Constants
MAX_LEN = 100  # Must match the training length
LABELS = {-1: "Negative", 0: "Neutral", 1: "Positive"}
# Map Keras output indices (0, 1, 2) to Sentiment Labels (-1, 0, 1)
INDEX_TO_LABEL = {0: -1, 1: 0, 2: 1}

# --- Prediction Logic ---

def predict_single(text, model_name):
    if not text: return "Please enter text."
    clean_text = preprocess(text)
    
    # A. Scikit-Learn Models
    if model_name in ["Logistic Regression", "Naive Bayes"]:
        model = lr_model if model_name == "Logistic Regression" else nb_model
        vectorized = tfidf.transform([clean_text])
        pred = model.predict(vectorized)[0]
        return LABELS.get(int(pred), "Unknown")

    # B. Feedforward NN (TF-IDF based)
    elif model_name == "Feedforward Neural Network":
        vectorized = tfidf.transform([clean_text]).toarray() # NN needs dense array
        pred_probs = ffnn_model.predict(vectorized)
        pred_idx = np.argmax(pred_probs, axis=1)[0]
        return LABELS.get(INDEX_TO_LABEL[pred_idx], "Unknown")

    # C. Embedding NN (Sequence based)
    elif model_name == "Embedding Neural Network":
        seq = tokenizer.texts_to_sequences([clean_text])
        padded = pad_sequences(seq, maxlen=MAX_LEN, padding='post', truncating='post')
        pred_probs = embed_model.predict(padded)
        pred_idx = np.argmax(pred_probs, axis=1)[0]
        return LABELS.get(INDEX_TO_LABEL[pred_idx], "Unknown")

def predict_file(file_obj, model_name):
    if file_obj is None: return None
    df = pd.read_csv(file_obj.name)
    
    if 'clean_text' not in df.columns:
        return "Error: CSV must have a column named 'clean_text'"
    
    # Preprocess all rows
    processed_texts = df['clean_text'].apply(preprocess)
    
    # A. Scikit-Learn Batch
    if model_name in ["Logistic Regression", "Naive Bayes"]:
        model = lr_model if model_name == "Logistic Regression" else nb_model
        vectors = tfidf.transform(processed_texts)
        preds = model.predict(vectors)
        df['Predicted_Sentiment'] = [LABELS.get(int(p), "Unknown") for p in preds]

    # B. Feedforward NN Batch
    elif model_name == "Feedforward Neural Network":
        # Note: Large datasets might need batching here to avoid memory issues with .toarray()
        vectors = tfidf.transform(processed_texts).toarray()
        pred_probs = ffnn_model.predict(vectors)
        pred_indices = np.argmax(pred_probs, axis=1)
        df['Predicted_Sentiment'] = [LABELS.get(INDEX_TO_LABEL[i], "Unknown") for i in pred_indices]

    # C. Embedding NN Batch
    elif model_name == "Embedding Neural Network":
        seqs = tokenizer.texts_to_sequences(processed_texts)
        padded = pad_sequences(seqs, maxlen=MAX_LEN, padding='post', truncating='post')
        pred_probs = embed_model.predict(padded)
        pred_indices = np.argmax(pred_probs, axis=1)
        df['Predicted_Sentiment'] = [LABELS.get(INDEX_TO_LABEL[i], "Unknown") for i in pred_indices]

    output_file = "predictions.csv"
    df.to_csv(output_file, index=False)
    return output_file

# --- Interface ---

model_choices = [
    "Logistic Regression", 
    "Naive Bayes", 
    "Feedforward Neural Network", 
    "Embedding Neural Network"
]

with gr.Blocks(title="Twitter Sentiment Analysis") as demo:
    gr.Markdown("# Twitter Sentiment Analysis Hub")
    
    with gr.Tabs():
        # Tab 1: Single Tweet
        with gr.TabItem("Analyze Tweet"):
            with gr.Row():
                txt_input = gr.Textbox(lines=2, placeholder="Enter tweet...", label="Input Text")
                model_drop = gr.Dropdown(model_choices, value="Logistic Regression", label="Select Model")
            btn_predict = gr.Button("Predict")
            lbl_output = gr.Label(label="Sentiment")
            
            btn_predict.click(predict_single, inputs=[txt_input, model_drop], outputs=lbl_output)

        # Tab 2: Batch Analysis
        with gr.TabItem("Batch Analysis (Upload CSV)"):
            gr.Markdown("Upload a CSV file with a `clean_text` column.")
            with gr.Row():
                file_input = gr.File(label="Upload CSV")
                model_drop_batch = gr.Dropdown(model_choices, value="Logistic Regression", label="Select Model")
            btn_process = gr.Button("Run Analysis")
            file_output = gr.File(label="Download Predictions")
            
            btn_process.click(predict_file, inputs=[file_input, model_drop_batch], outputs=file_output)

if __name__ == "__main__":
    demo.launch()