# -*- coding: utf-8 -*-
"""27to99RNN.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1c1amKG3mOiR3kr3q3HZX01pm7gkRWUkw
"""

# Commented out IPython magic to ensure Python compatibility.
# -*- coding: utf-8 -*-
"""Enhanced_NLP_PROJECT_rnn.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1ZdzAVjrcHeK1QpehyRV85JD_FzNdOifg
"""

# ----------------------------- #
#        1. Environment Setup   #
# ----------------------------- #

# Install necessary packages
!pip install kaggle
!pip install librosa
!pip install soundfile

# Import necessary libraries
import os
import zipfile
import numpy as np
import pandas as pd
import librosa
import librosa.display
import soundfile as sf
from sklearn.preprocessing import LabelEncoder, OneHotEncoder
from sklearn.utils import class_weight
from sklearn.model_selection import train_test_split
from keras.utils import to_categorical
import matplotlib.pyplot as plt
import seaborn as sns
import warnings
warnings.filterwarnings('ignore')

# Set matplotlib to display inline
# %matplotlib inline

# Set random seed for reproducibility
import tensorflow as tf
import random
seed = 42
random.seed(seed)
np.random.seed(seed)
tf.random.set_seed(seed)

# ----------------------------- #
#        2. Data Download        #
# ----------------------------- #

# Ensure you have your Kaggle API token in /content
# You can upload the kaggle.json file by clicking the folder icon on the left sidebar in Colab

# Set Kaggle configuration directory
os.environ['KAGGLE_CONFIG_DIR'] = "/content"

# Download the TESS dataset from Kaggle
!kaggle datasets download -d ejlok1/toronto-emotional-speech-set-tess

# Unzip the dataset
with zipfile.ZipFile("toronto-emotional-speech-set-tess.zip", 'r') as zip_ref:
    zip_ref.extractall()

# Remove the zip file to save space
os.remove("toronto-emotional-speech-set-tess.zip")

# ----------------------------- #
#        3. Data Loading         #
# ----------------------------- #

# Define the path to the dataset
DATASET_PATH = '/content/TESS Toronto emotional speech set data'

# Initialize lists to hold file paths and labels
paths = []
labels = []

# Traverse the dataset directory to collect file paths and labels
for dirname, _, filenames in os.walk(DATASET_PATH):
    for filename in filenames:
        paths.append(os.path.join(dirname, filename))
        label = filename.split('_')[-1]
        label = label.split('.')[0]
        labels.append(label.lower())
    if len(paths) >= 2800:
        break  # Limit to 2800 samples for manageability

print('Dataset is Loaded')

# Display dataset statistics
print(f'Number of samples: {len(paths)}')
print('Sample paths:', paths[:5])
print('Sample labels:', labels[:5])

# Create a DataFrame
df = pd.DataFrame({'speech': paths, 'label': labels})
print(df.head())
print(df['label'].value_counts())

# ----------------------------- #
#        4. Exploratory Data Analysis (EDA)  #
# ----------------------------- #

# Plot the distribution of emotions
plt.figure(figsize=(10,6))
sns.countplot(data=df, x='label')
plt.title('Distribution of Emotion Classes')
plt.xlabel('Emotion')
plt.ylabel('Count')
plt.show()

# Define utility functions to visualize waveforms and spectrograms
def waveplot(data, sr, emotion):
    plt.figure(figsize=(10,4))
    plt.title(f'Waveform - {emotion}', size=20)
    librosa.display.waveshow(data, sr=sr)
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    plt.show()

def spectrogram_plot(data, sr, emotion):
    X = librosa.stft(data)
    Xdb = librosa.amplitude_to_db(abs(X))
    plt.figure(figsize=(11,4))
    plt.title(f'Spectrogram - {emotion}', size=20)
    librosa.display.specshow(Xdb, sr=sr, x_axis='time', y_axis='hz')
    plt.colorbar(format='%+2.0f dB')
    plt.show()

# Display sample audio waves and spectrograms for each emotion
emotions = df['label'].unique()
for emotion in emotions:
    path = df['speech'][df['label'] == emotion].iloc[0]
    data, sampling_rate = librosa.load(path)
    waveplot(data, sampling_rate, emotion)
    spectrogram_plot(data, sampling_rate, emotion)
    # Uncomment the line below to listen to the audio (not recommended for large datasets)
    # display(Audio(path))

# ----------------------------- #
#        5. Data Augmentation    #
# ----------------------------- #

# Define a function to apply various augmentations to an audio signal
def augment_audio(y, sr):
    augmentations = []

    # 1. Additive Gaussian Noise
    noise = np.random.randn(len(y))
    augmented_y = y + 0.005 * noise
    augmentations.append(augmented_y)

    # 2. Time Shifting
    shift = np.random.randint(int(sr * 0.2))  # Shift by up to 20% of the length
    augmented_y = np.roll(y, shift)
    augmentations.append(augmented_y)

    # 3. Pitch Shifting
    augmented_y = librosa.effects.pitch_shift(y, sr, n_steps=2)  # Shift by +2 steps
    augmentations.append(augmented_y)

    # 4. Time Stretching
    augmented_y = librosa.effects.time_stretch(y, rate=1.1)  # Stretch by a factor of 1.1
    # To ensure the stretched signal has the same length
    if len(augmented_y) > len(y):
        augmented_y = augmented_y[:len(y)]
    else:
        augmented_y = np.pad(augmented_y, (0, len(y) - len(augmented_y)), 'constant')
    augmentations.append(augmented_y)

    return augmentations

# Apply augmentation and expand the dataset
augmented_paths = []
augmented_labels = []

print("Applying data augmentation...")
for index, row in df.iterrows():
    try:
        y, sr = librosa.load(row['speech'], duration=3, offset=0.5)
        augmentations = augment_audio(y, sr)
        for augmented_y in augmentations:
            augmented_paths.append(augmented_y)
            augmented_labels.append(row['label'])
    except Exception as e:
        print(f"Error processing {row['speech']}: {e}")

print(f"Original dataset size: {df.shape}")
print(f"Augmented data size: {len(augmented_paths)}")

# Create DataFrame for augmented data
aug_df = pd.DataFrame({'speech': augmented_paths, 'label': augmented_labels})

# Combine original and augmented data
# Note: Original data loaded as file paths, augmented data loaded as numpy arrays
# To maintain consistency, we will extract features from augmented numpy arrays directly
print("Combining original and augmented data...")
combined_speeches = df['speech'].tolist() + aug_df['speech'].tolist()
combined_labels = df['label'].tolist() + aug_df['label'].tolist()

print(f"Combined dataset size: {len(combined_speeches)}")

# ----------------------------- #
#        6. Feature Extraction   #
# ----------------------------- #

# Define a function to extract audio features
def extract_features(y, sr):
    # Extract MFCC
    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=40)

    # Extract Chroma
    stft = np.abs(librosa.stft(y))
    chroma = librosa.feature.chroma_stft(S=stft, sr=sr)

    # Extract Spectral Contrast
    contrast = librosa.feature.spectral_contrast(S=stft, sr=sr)

    # Extract Tonnetz
    y_harmonic = librosa.effects.harmonic(y)
    tonnetz = librosa.feature.tonnetz(y=y_harmonic, sr=sr)

    # Combine all features
    features = np.vstack([mfcc, chroma, contrast, tonnetz])

    return features.T  # Shape: (time_steps, features)

# Extract features for all samples
print("Extracting features from all samples...")
X_features = []
for i, speech in enumerate(combined_speeches):
    if isinstance(speech, str):
        # Original data: load from file
        y, sr = librosa.load(speech, duration=3, offset=0.5)
    else:
        # Augmented data: already loaded as numpy array
        y = speech
        sr = 22050  # Default sampling rate
    features = extract_features(y, sr)
    X_features.append(features)
    if (i+1) % 500 == 0:
        print(f"Processed {i+1}/{len(combined_speeches)} samples")

print("Feature extraction completed.")

# Determine the maximum sequence length
max_len = max([x.shape[0] for x in X_features])
print(f'Max sequence length: {max_len}')

# Pad sequences to have the same length
from keras.preprocessing.sequence import pad_sequences

print("Padding sequences...")
X_padded = pad_sequences(X_features, maxlen=max_len, dtype='float32', padding='post', truncating='post')

# Expand dimensions for LSTM input (samples, time_steps, features, 1)
# For Conv1D + LSTM, input shape should be (samples, time_steps, features)
X = X_padded  # Shape: (samples, time_steps, features)

print('Padded features shape:', X.shape)

# ----------------------------- #
#        7. Label Encoding       #
# ----------------------------- #

# Encode labels as integers
print("Encoding labels...")
le = LabelEncoder()
y_int = le.fit_transform(combined_labels)

# One-hot encode labels
y = to_categorical(y_int, num_classes=len(le.classes_))
print('One-hot encoded labels shape:', y.shape)

# ----------------------------- #
#        8. Handle Class Imbalance #
# ----------------------------- #

# Compute class weights to handle imbalance
print("Computing class weights...")
class_weights_array = class_weight.compute_class_weight('balanced', classes=np.unique(y_int), y=y_int)
# this is extra line below 2----------------------------------------------------------------------------------------------
class_weights = dict(enumerate(class_weights_array))
print(f'Class weights: {class_weights}')

# ----------------------------- #
#        9. Train-Test Split     #
# ----------------------------- #

# Split the data into training and testing sets
print("Splitting data into training and testing sets...")
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=seed, stratify=y)

print(f'Training set size: {X_train.shape}')
print(f'Testing set size: {X_test.shape}')

# ----------------------------- #
#        10. Model Definition    #
# ----------------------------- #

from keras.models import Sequential
from keras.layers import Dense, LSTM, Dropout, Bidirectional, Conv1D, MaxPooling1D, BatchNormalization, GlobalMaxPooling1D

# Define the model architecture
print("Building the model...")

model = Sequential()

# 1. Convolutional Layer
model.add(Conv1D(filters=64, kernel_size=3, activation='relu', input_shape=(X_train.shape[1], X_train.shape[2])))
model.add(BatchNormalization())
model.add(MaxPooling1D(pool_size=2))
model.add(Dropout(0.3))

# 2. Bidirectional LSTM Layer
model.add(Bidirectional(LSTM(128, return_sequences=True)))
model.add(BatchNormalization())
model.add(Dropout(0.3))

# 3. Another Bidirectional LSTM Layer
model.add(Bidirectional(LSTM(64)))
model.add(BatchNormalization())
model.add(Dropout(0.3))

# 4. Fully Connected Layers
model.add(Dense(128, activation='relu'))
model.add(Dropout(0.3))
model.add(Dense(64, activation='relu'))
model.add(Dropout(0.3))

# 5. Output Layer
model.add(Dense(len(le.classes_), activation='softmax'))

# Compile the model
from keras.optimizers import Adam

optimizer = Adam(learning_rate=0.0001)
model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])
model.summary()

# ----------------------------- #
#        11. Callbacks Setup    #
# ----------------------------- #

from keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau

# Define callbacks
early_stop = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True, verbose=1)
# Change the file extension to .keras
checkpoint = ModelCheckpoint('best_model.keras', monitor='val_loss', save_best_only=True, verbose=1)
reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5, min_lr=0.00001, verbose=1)

# ----------------------------- #
#        12. Model Training      #
# ----------------------------- #

# Train the model
print("Starting training...")
history = model.fit(X_train, y_train,
                    validation_split=0.2,
                    epochs=100,
                    batch_size=64,
                    class_weight=class_weights,
                    callbacks=[early_stop, checkpoint, reduce_lr],
                    verbose=1)

# ----------------------------- #
#        13. Plot Training History #
# ----------------------------- #

# Plot Accuracy
plt.figure(figsize=(12, 6))
plt.plot(history.history['accuracy'], label='Train Accuracy')
plt.plot(history.history['val_accuracy'], label='Validation Accuracy')
plt.title('Accuracy Over Epochs')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

# Plot Loss
plt.figure(figsize=(12, 6))
plt.plot(history.history['loss'], label='Train Loss')
plt.plot(history.history['val_loss'], label='Validation Loss')
plt.title('Loss Over Epochs')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()

# ----------------------------- #
#        14. Load Best Model     #
##        14. Load Best Model     #
# ----------------------------- #

# Load the best model saved during training
print("Loading the best model...")
from keras.models import load_model
# Change the file name to 'best_model.keras'
best_model = load_model('best_model.keras')

# ----------------------------- #
#        15. Evaluate the Model  #
# ----------------------------- #

# Evaluate on the test set
print("Evaluating the model on the test set...")
test_loss, test_acc = best_model.evaluate(X_test, y_test, verbose=0)
print(f'Test Accuracy: {test_acc * 100:.2f}%')
print(f'Test Loss: {test_loss:.4f}')

# ----------------------------- #
#        16. Save the Model      #
# ----------------------------- #

# Save the model for future use
print("Saving the model...")
best_model.save('/content/enhanced_model_rnn.h5')

# ----------------------------- #
#        17. Prediction Function #
# ----------------------------- #

# Define a function to preprocess and extract features from a new audio file
def preprocess_audio(file_path):
    try:
        y, sr = librosa.load(file_path, duration=3, offset=0.5)
        features = extract_features(y, sr)
        # Pad the sequence
        features_padded = pad_sequences([features], maxlen=max_len, dtype='float32', padding='post', truncating='post')
        return features_padded
    except Exception as e:
        print(f"Error processing {file_path}: {e}")
        return None

# Define a function to predict emotion from a new audio file
def predict_emotion(file_path, model):
    features = preprocess_audio(file_path)
    if features is not None:
        prediction = model.predict(features)
        predicted_label_index = np.argmax(prediction, axis=1)[0]
        predicted_emotion = le.inverse_transform([predicted_label_index])[0]
        return predicted_emotion
    else:
        return "Error in processing audio."

# ----------------------------- #
#        18. Test Prediction     #
# ----------------------------- #

# Example prediction with a new audio file
# Ensure you have the MP3 file uploaded to Colab or accessible via a file path
new_file_path = '/content/zzz01j79-fm6wj-5mw42-3bbzz-x45618.mp3'  # Replace with your MP3 file path

# If the file is not already present, upload it
from google.colab import files
import shutil

if not os.path.exists(new_file_path):
    print("Please upload the new audio file.")
    uploaded = files.upload()
    for filename in uploaded.keys():
        shutil.move(filename, new_file_path)
        print(f"Uploaded and moved {filename} to {new_file_path}")

# Predict the emotion
predicted_emotion = predict_emotion(new_file_path, best_model)
print(f'The predicted emotion for the new file is: {predicted_emotion}')

print(f'Test Accuracy: {test_acc * 100:.2f}%')

in this code print graph of the result like f1 score presicio accuracy all graphs