import numpy as np
import librosa.display
import matplotlib.pyplot as plt
import os

# rename files in ./data/*

def create_spectrogram(audio_file, image_file):
    y, sr = librosa.load(audio_file)
    plt.figure(figsize=(12, 8))
    D = librosa.amplitude_to_db(np.abs(librosa.stft(y)), ref=np.max)
    librosa.display.specshow(D, y_axis='linear')
    plt.colorbar(format='%+2.0f dB')
    plt.savefig(image_file, bbox_inches='tight', pad_inches=0.0)
    plt.close()

def create_pngs_from_wavs(input_path, output_path):
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    dir = os.listdir(input_path)

    for i, file in enumerate(dir):
        input_file = os.path.join(input_path, file)
        output_file = os.path.join(output_path, file.replace('.wav', '.png'))
        create_spectrogram(input_file, output_file)


create_pngs_from_wavs('./data/a/', './data/a_png/')
create_pngs_from_wavs('./data/b/', './data/b_png/')
create_pngs_from_wavs('./data/c/', './data/c_png/')