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

Automatically generated by Colab.

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

import gradio as gr
import numpy as np
import librosa
import os
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import warnings

warnings.filterwarnings("ignore")

# === Setup: Load and extract training features ===
def extract_features(file_path):
    try:
        audio, sr = librosa.load(file_path, duration=3, offset=0.5)
        mfccs = librosa.feature.mfcc(y=audio, sr=sr, n_mfcc=13)
        return np.mean(mfccs.T, axis=0)
    except Exception as e:
        print("Error:", e)
        return None

# Folder where your dataset is stored (update if needed)
DATA_DIR = "/content/drive/MyDrive/Audio"

# Load and prepare dataset
features, labels = [], []
model = None # Initialize model to None
if os.path.exists(DATA_DIR):
    for folder in os.listdir(DATA_DIR):
        emotion = folder.split('_')[-1].lower()
        folder_path = os.path.join(DATA_DIR, folder)
        if os.path.isdir(folder_path): # Check if it's a directory
            for file in os.listdir(folder_path):
                if file.endswith(".wav"):
                    file_path = os.path.join(folder_path, file)
                    mfcc = extract_features(file_path)
                    if mfcc is not None:
                        features.append(mfcc)
                        labels.append(emotion)

    if features and labels: # Check if data was loaded
        X = np.array(features)
        y = np.array(labels)

        # Train model
        model = RandomForestClassifier()
        model.fit(X, y)
        print("Model trained successfully.")
    else:
        print(f"No audio files found in '{DATA_DIR}' or subfolders.")
else:
    print(f"ERROR: Folder '{DATA_DIR}' not found. Please upload your training data folder.")


# === Prediction function ===
def predict_emotion(audio_file):
    if model is None:
        return "Model not trained. Upload audio_dataset folder with training data and re-run the code."

    features = extract_features(audio_file)
    if features is not None:
        try:
            prediction = model.predict(features.reshape(1, -1))[0]
            return f"Predicted Emotion: {prediction}"
        except Exception as e:
            return f"Prediction error: {e}"
    else:
        return "Could not process the audio."

# === Gradio App ===
gr.Interface(
    fn=predict_emotion,
    inputs=gr.Audio(type="filepath", label="Upload a .wav file"), # Removed source parameter
    outputs="text",
    title="Speech Emotion Recognition",
    description="Upload a 3-second .wav file to predict the speaker's emotion.\n\nNote: You must have a folder named 'audio_dataset' structured like 'speaker_happy/filename.wav'."
).launch()