import os
import librosa
import librosa.display
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm  # <-- import tqdm

def audio_to_spectrogram(audio_path, save_path):
    y, sr = librosa.load(audio_path, sr=None)
    plt.figure(figsize=(3, 3))
    D = librosa.amplitude_to_db(np.abs(librosa.stft(y)), ref=np.max)
    librosa.display.specshow(D, sr=sr, x_axis='time', y_axis='log')
    plt.axis('off')  # to not display axes

    # Save only the image without any white spaces or borders
    plt.tight_layout(pad=0)
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
    plt.savefig(save_path, bbox_inches='tight', pad_inches=0, transparent=True)
    plt.close()

def process_audio_directory(source_folder, dest_folder, is_test=False):
    for root, dirs, files in os.walk(source_folder):
        if not is_test:
            for dir_name in dirs:
                # Create a corresponding directory in the destination folder if it doesn't exist
                os.makedirs(os.path.join(dest_folder, dir_name), exist_ok=True)
        
        for file in tqdm(files, desc="Processing files in {}".format(os.path.basename(root)), unit="file"):
            if file.endswith('.wav'):
                if is_test:
                    audio_path = os.path.join(root, file)
                    image_name = file.replace('.wav', '.png')
                    image_path = os.path.join(dest_folder, image_name)
                else:
                    class_name = os.path.basename(root)
                    audio_path = os.path.join(root, file)
                    image_name = file.replace('.wav', '.png')
                    image_path = os.path.join(dest_folder, class_name, image_name)
                audio_to_spectrogram(audio_path, image_path)

source_train_audio = 'audio_dataset/train'
dest_train_images = 'image_dataset/train'
print("Processing Train Dataset:")
process_audio_directory(source_train_audio, dest_train_images)

source_test_audio = 'audio_dataset/test'
dest_test_images = 'image_dataset/test'
print("\nProcessing Test Dataset:")
process_audio_directory(source_test_audio, dest_test_images, is_test=True)