===============================
# 1. Import Libraries
# ===============================
import os
import numpy as np
import pandas as pd
import librosa
import tensorflow as tf
from tensorflow.keras import layers, models
from sklearn.model_selection import train_test_split

# ===============================
# 2. Dataset Path (YOUR CORRECT PATH)
# ===============================
DATASET_PATH = r"C:\Users\kaviy\Downloads\kaviyadataset\urbansound8k"

csv_path = os.path.join(DATASET_PATH, "metadata", "UrbanSound8K.csv")
audio_path = os.path.join(DATASET_PATH, "audio")

# Check path
print("CSV Exists:", os.path.exists(csv_path))

# Load metadata
metadata = pd.read_csv(csv_path)

# Create full audio path
metadata['path'] = metadata.apply(
    lambda row: os.path.join(audio_path, f"fold{row['fold']}", row['slice_file_name']),
    axis=1
)

print("Metadata Loaded ✅")

# ===============================
# 3. Feature Extraction (Mel Spectrogram)
# ===============================
def extract_features(file_path, max_len=128):
    try:
        audio, sr = librosa.load(file_path, sr=22050)

        mel = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=128)
        mel_db = librosa.power_to_db(mel, ref=np.max)

        # Fix size
        if mel_db.shape[1] < max_len:
            pad = max_len - mel_db.shape[1]
            mel_db = np.pad(mel_db, ((0,0),(0,pad)), mode='constant')
        else:
            mel_db = mel_db[:, :max_len]

        return mel_db

    except:
        return None


# ===============================
# 4. Prepare Dataset
# ===============================
X, y = [], []

for i, row in metadata.iterrows():
    feature = extract_features(row['path'])

    if feature is not None:
        X.append(feature)
        y.append(row['classID'])

X = np.array(X)
y = np.array(y)

print("Dataset Shape:", X.shape)

# Add channel dimension
X = X[..., np.newaxis]

# Convert to 3 channels (for CNN)
X = np.repeat(X, 3, axis=-1)

# ===============================
# 5. Train-Test Split
# ===============================
X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,
    stratify=y,
    random_state=42
)

# ===============================
# 6. Transfer Learning Model
# ===============================
base_model = tf.keras.applications.MobileNetV2(
    weights='imagenet',
    include_top=False,
    input_shape=(128,128,3)
)

base_model.trainable = False

model = models.Sequential([
    base_model,
    layers.GlobalAveragePooling2D(),
    layers.Dense(128, activation='relu'),
    layers.Dropout(0.3),
    layers.Dense(10, activation='softmax')
])

model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

model.summary()

# ===============================
# 7. Training
# ===============================
history = model.fit(
    X_train,
    y_train,
    validation_data=(X_test, y_test),
    epochs=5,
    batch_size=16
)

# ===============================
# 8. Evaluation
# ===============================
loss, acc = model.evaluate(X_test, y_test)
print("Test Accuracy:", acc)

# ===============================
# 9. Prediction
# ===============================
pred = model.predict(X_test[:1])
print("Predicted Class:", np.argmax(pred))



