import tensorflow as tf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences

# Ensure TensorFlow uses GPU if available
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if physical_devices:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)
    print("GPU is available and will be used.")
else:
    print("No GPU found. Using CPU.")

# Load IMDb dataset
df = pd.read_csv("dataset/imdb.csv")

# Check column names
print("Dataset columns:", df.columns)

# Rename columns if needed
df.rename(columns={'review': 'text', 'sentiment': 'label'}, inplace=True)

# Convert labels to binary (positive = 1, negative = 0)
df['label'] = df['label'].apply(lambda x: 1 if x.strip().lower() == 'positive' else 0)

# Tokenize text
tokenizer = Tokenizer(num_words=20000, oov_token="<OOV>")  # Increased vocab size
tokenizer.fit_on_texts(df['text'])
sequences = tokenizer.texts_to_sequences(df['text'])
padded_sequences = pad_sequences(sequences, maxlen=250)  # Increased max length

# Split data
X_train, y_train = padded_sequences[:20000], df['label'][:20000]
X_test, y_test = padded_sequences[20000:], df['label'][20000:]

# Build a more complex neural network model
model = tf.keras.Sequential([
    tf.keras.layers.Embedding(20000, 64, input_length=250),  # Larger embedding space
    tf.keras.layers.Conv1D(128, 5, activation='relu'),
    tf.keras.layers.MaxPooling1D(pool_size=2),
    tf.keras.layers.LSTM(128, return_sequences=True),
    tf.keras.layers.LSTM(64),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

# Compile model
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# Train model epochs
epochs = 5
history = model.fit(X_train, y_train, epochs=epochs, validation_data=(X_test, y_test), batch_size=128)

# Save the final trained model
model.save("src/sentiment_model_final.h5")

# Save intermediate models every 10 epochs
for i in range(10, epochs + 1, 10):
    model.save(f"src/sentiment_model_epoch_{i}.h5")

# Plot detailed training loss and accuracy with smoothing
def smooth_curve(points, factor=0.8):
    """Applies exponential moving average smoothing to a curve."""
    smoothed_points = []
    for point in points:
        if smoothed_points:
            smoothed_points.append(smoothed_points[-1] * factor + point * (1 - factor))
        else:
            smoothed_points.append(point)
    return smoothed_points

plt.figure(figsize=(12, 6))

# Loss Plot
plt.subplot(1, 2, 1)
plt.plot(range(epochs), history.history['loss'], label='Training Loss', marker='o', alpha=0.3)
plt.plot(range(epochs), history.history['val_loss'], label='Validation Loss', marker='o', alpha=0.3)
plt.plot(range(epochs), smooth_curve(history.history['loss']), label='Smoothed Training Loss', linewidth=2)
plt.plot(range(epochs), smooth_curve(history.history['val_loss']), label='Smoothed Validation Loss', linewidth=2)
plt.xlabel("Epochs")
plt.ylabel("Loss")
plt.legend()
plt.title("Training vs. Validation Loss (Detailed)")

# Accuracy Plot
plt.subplot(1, 2, 2)
plt.plot(range(epochs), history.history['accuracy'], label='Training Accuracy', marker='o', alpha=0.3)
plt.plot(range(epochs), history.history['val_accuracy'], label='Validation Accuracy', marker='o', alpha=0.3)
plt.plot(range(epochs), smooth_curve(history.history['accuracy']), label='Smoothed Training Accuracy', linewidth=2)
plt.plot(range(epochs), smooth_curve(history.history['val_accuracy']), label='Smoothed Validation Accuracy', linewidth=2)
plt.xlabel("Epochs")
plt.ylabel("Accuracy")
plt.legend()
plt.title("Training vs. Validation Accuracy (Detailed)")

# Save and show graph
plt.savefig("src/training_results_detailed.png")
plt.show()