# -*- coding: utf-8 -*-
"""Untitled11.ipynb

Automatically generated by Colab.

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

# -*- coding: utf-8 -*-
"""96%cnnECHONLPipynb"""

#dataset
import os
os.environ['KAGGLE_CONFIG_DIR'] = "/content"

!kaggle datasets download -d ejlok1/toronto-emotional-speech-set-tess
!unzip \*.zip  && rm *.zip

# Import necessary modules
import pandas as pd
import numpy as np
import os
import seaborn as sns
import matplotlib.pyplot as plt
import librosa
import librosa.display
from IPython.display import Audio
import warnings
warnings.filterwarnings('ignore')

# Load the dataset
paths = []
labels = []
for dirname, _, filenames in os.walk('/content/TESS Toronto emotional speech set data'):
    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
print('Dataset is Loaded')

# Check dataset length and sample paths
print(f'Number of samples: {len(paths)}')
print('Sample paths:', paths[:5])
print('Sample labels:', labels[:5])

# Create a dataframe
df = pd.DataFrame()
df['audio_paths'] = paths
df['labels'] = labels
print(df.head())
print(df['labels'].value_counts())

# Exploratory Data Analysis
sns.countplot(data=df, x='labels')
plt.show()

# Function to plot the waveform with time and amplitude labels
def waveplot(data, sr, emotion):
    plt.figure(figsize=(10, 4))
    plt.title(f'Waveform of {emotion.capitalize()}', size=16, color='navy')
    librosa.display.waveshow(data, sr=sr, alpha=0.8, color='mediumseagreen')

    # Adding time axis
    plt.xlabel("Time (seconds)", fontsize=12)

    # Adding amplitude axis
    plt.ylabel("Amplitude", fontsize=12)

    plt.grid(True, linestyle='--', alpha=0.7)
    plt.show()

# Function to plot a unique spectrogram with both time and frequency axis and detailed color bar
def spectogram(data, sr, emotion):
    stft = librosa.stft(data, n_fft=2048, hop_length=512)
    stft_db = librosa.amplitude_to_db(abs(stft))

    plt.figure(figsize=(10, 4))
    plt.title(f'Spectrogram of {emotion.capitalize()}', size=16, color='darkred')

    # Display the spectrogram with a unique color map
    librosa.display.specshow(stft_db, sr=sr, hop_length=512, x_axis='time', y_axis='log', cmap='plasma')

    # Add color bar with labels in decibels (dB)
    cbar = plt.colorbar(format="%+2.0f dB")
    cbar.set_label('Amplitude (dB)', fontsize=12)

    # Add time and frequency axis labels
    plt.xlabel("Time (seconds)", fontsize=12)
    plt.ylabel("Frequency (Hz)", fontsize=12)

    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()

# Display sample audio waves and spectrograms with unique styling
emotions = ['fear', 'angry', 'disgust', 'neutral', 'sad', 'ps', 'happy']
for emotion in emotions:
    path = np.array(df['audio_paths'][df['labels'] == emotion])[0]
    data, sampling_rate = librosa.load(path, sr=16000)

    # Plot waveform
    waveplot(data, sampling_rate, emotion)

    # Plot spectrogram
    spectogram(data, sampling_rate, emotion)

    # Play the audio sample
    display(Audio(path))
# Create label to integer mapping
label_map = {label: idx for idx, label in enumerate(df['labels'].unique())}
inverse_label_map = {idx: label for label, idx in label_map.items()}

print('Label to Integer Mapping:', label_map)
print('Integer to Label Mapping:', inverse_label_map)

# Convert labels in the dataframe to integers
df['labels'] = df['labels'].map(label_map)
print(df.head(2))



import torch
import torchaudio
from torch.utils.data import Dataset

class SpeechEmotionDataset(Dataset):
    def __init__(self, df, processor, max_length=32000):
        self.df = df
        self.processor = processor
        self.max_length = max_length

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        audio_path = self.df.iloc[idx]['audio_paths']
        label = self.df.iloc[idx]['labels']  # This is already an integer

        # Use librosa for loading and resampling
        speech, sr = librosa.load(audio_path, sr=16000)

        # Pad or truncate the speech to the required length
        if len(speech) > self.max_length:
            speech = speech[:self.max_length]
        else:
            speech = np.pad(speech, (0, self.max_length - len(speech)), 'constant')

        # Preprocess the audio file
        inputs = self.processor(speech, sampling_rate=16000, return_tensors="pt", padding=True, truncation=True, max_length=self.max_length)
        input_values = inputs.input_values.squeeze()

        return {'input_values': input_values, 'labels': torch.tensor(label, dtype=torch.long)}  # Label is already integer


# Split the DataFrame into training and testing sets
from sklearn.model_selection import train_test_split
train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)

from transformers import Wav2Vec2Processor, Wav2Vec2ForSequenceClassification

# Load the Wav2Vec2 processor and model for sequence classification
processor = Wav2Vec2Processor.from_pretrained('facebook/wav2vec2-base')
model = Wav2Vec2ForSequenceClassification.from_pretrained('facebook/wav2vec2-base', num_labels=len(label_map))
train_dataset = SpeechEmotionDataset(train_df, processor, max_length=32000) # Pass max_length explicitly
test_dataset = SpeechEmotionDataset(test_df, processor, max_length=32000) # Pass max_length explicitly


# Create DataLoader for training and testing datasets
from torch.utils.data import DataLoader
train_dataloader = DataLoader(train_dataset, batch_size=8, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=8, shuffle=False)

# Training pipeline setup
from transformers import Trainer, TrainingArguments
from sklearn.metrics import accuracy_score, precision_recall_fscore_support

# Define compute metrics function for accuracy and F1 score
def compute_metrics(pred):
    labels = pred.label_ids
    preds = np.argmax(pred.predictions, axis=1)
    accuracy = accuracy_score(labels, preds)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='weighted')

    return {
        'accuracy': accuracy,
        'precision': precision,
        'recall': recall,
        'f1': f1
    }

training_args = TrainingArguments(
    output_dir='./results',
    evaluation_strategy='epoch',
    save_strategy='epoch',
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    weight_decay=0.01,
    report_to=[],
)

# Initialize the Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    compute_metrics=compute_metrics
)

# Start training
trainer.train()

# Evaluate the model and print results
results = trainer.evaluate()
print(results)

# Save the model's state_dict instead of the entire model
torch.save(model.state_dict(), 'cnn_model_full_final.pth')

# Test Prediction
import torchaudio
import numpy as np
import torch

# Preprocess audio
def preprocess_audio(file_path, processor, max_length=32000):
    audio, sr = torchaudio.load(file_path)
    audio = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000)(audio).squeeze().numpy()
    audio = audio[:max_length] if len(audio) > max_length else np.pad(audio, (0, max_length - len(audio)), 'constant')
    return processor(audio, sampling_rate=16000, return_tensors="pt", padding=True, truncation=True, max_length=max_length).input_values.squeeze()

# Predict emotion
def predict_emotion(file_path, model, processor, label_map):
    inputs = preprocess_audio(file_path, processor)
    logits = model(inputs.unsqueeze(0)).logits
    predicted_class_id = torch.argmax(logits, dim=-1).item()
    return {v: k for k, v in label_map.items()}[predicted_class_id]

from google.colab import files

# Upload file
uploaded = files.upload()
file_path = list(uploaded.keys())[0]

# Prediction
predicted_emotion = predict_emotion(file_path, model, processor, label_map)
print(f'The predicted emotion is: {predicted_emotion}')in this code i want graph of model of the accuracy precision and etc of the model don't print the table of these instead i want 2 dimension of each