# Important: import tf-keras before transformers to avoid Keras 3 issue
import pandas as pd
import tensorflow as tf
import tf_keras as keras  # Import tf-keras before transformers to avoid Keras 3 issues

from transformers import T5Tokenizer, TFT5ForConditionalGeneration
from sklearn.model_selection import train_test_split
import gradio as gr


# ========== STEP 1: Load and Prepare Dataset ==========
# NOTE: On Hugging Face Spaces, upload your CSV file in the repo and update path here:
df = pd.read_csv("https://huggingface.co/datasets/gmustafa413/ABSA_Dataset/raw/main/Laptops_Train.csv")  # Make sure this CSV is uploaded in your repo root

# Convert 'aspectTerms' string to list and create input/output pairs
def format_example(row):
    input_text = row["raw_text"]
    try:
        aspect_terms = eval(row["aspectTerms"])
    except:
        aspect_terms = []

    if not aspect_terms or aspect_terms == '[]':
        output_text = "noaspectterm:none"
    else:
        output_text = ", ".join([f"{d['term']}:{d['polarity']}" for d in aspect_terms if d['term']])
    return input_text, output_text

df["input_text"], df["target_text"] = zip(*df.apply(format_example, axis=1))

# Split data
train_texts, val_texts, train_labels, val_labels = train_test_split(
    df["input_text"], df["target_text"], test_size=0.1, random_state=42
)

# ========== STEP 2: Tokenization ==========
tokenizer = T5Tokenizer.from_pretrained("t5-base")
max_input_length = 128
max_target_length = 64

def tokenize_data(inputs, targets):
    inputs_enc = tokenizer(
        list(inputs), padding="max_length", truncation=True,
        max_length=max_input_length, return_tensors="tf"
    )
    targets_enc = tokenizer(
        list(targets), padding="max_length", truncation=True,
        max_length=max_target_length, return_tensors="tf"
    )
    return inputs_enc.input_ids, targets_enc.input_ids

train_inputs, train_labels = tokenize_data(train_texts, train_labels)
val_inputs, val_labels = tokenize_data(val_texts, val_labels)

# ========== STEP 3: Create TensorFlow Datasets ==========
BATCH_SIZE = 8

train_dataset = tf.data.Dataset.from_tensor_slices(({"input_ids": train_inputs}, train_labels))
val_dataset = tf.data.Dataset.from_tensor_slices(({"input_ids": val_inputs}, val_labels))

train_dataset = train_dataset.shuffle(1000).batch(BATCH_SIZE)
val_dataset = val_dataset.batch(BATCH_SIZE)

# ========== STEP 4: Model Training ==========
model = TFT5ForConditionalGeneration.from_pretrained("t5-base")
optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5)
model.compile(optimizer=optimizer)

print("🔧 Training started...")
model.fit(train_dataset, validation_data=val_dataset, epochs=3)
print("✅ Training complete!")

# Save model and tokenizer (optional here if you want to save locally)
# model.save_pretrained("t5_absa_tf_model")
# tokenizer.save_pretrained("t5_absa_tf_model")

# ========== STEP 5: Gradio Interface ==========
# Using trained model & tokenizer directly from memory

def preprocess_input(text):
    instruction = """Definition: The output will be the aspects (both implicit and explicit) and the aspect's sentiment polarity. In cases where there are no aspects the output should be noaspectterm:none.
Positive example:
input: With the great variety on the menu , I eat here often and never get bored.
output: menu:positive
Now complete the following example-
input: """
    return instruction + text + "\noutput:"

def predict_aspects(text):
    processed = preprocess_input(text)
    inputs = tokenizer(processed, return_tensors="tf", truncation=True, padding=True)
    outputs = model.generate(
        input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"],
        max_length=64
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

interface = gr.Interface(
    fn=predict_aspects,
    inputs=gr.Textbox(lines=2, placeholder="Enter your review here..."),
    outputs=gr.Textbox(label="Extracted aspect:sentiment pairs"),
    title="Aspect-Based Sentiment Analysis (ABSA)",
    description="Fine-tuned T5 model for extracting aspects and their sentiment from review text.",
)

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