# standard import 
import glob
import argparse
import sys
import os

# additional import 
import numpy
import librosa
import librosa.core
import librosa.feature
import yaml
from tqdm import tqdm


########################################################################
# Extract all files ending with .wav from a directory 
########################################################################
def file_list_generator(target_dir,
                        dir_name="train",
                        ext="wav"):

    # generate training list
    training_list_path = os.path.abspath("{dir}/{dir_name}/*.{ext}".format(dir=target_dir, dir_name=dir_name, ext=ext))
    files = sorted(glob.glob(training_list_path))

    print("Train file num : {num}".format(num=len(files)))
    return files

########################################################################
# Extract all files ending with .wav from a directory for testing
########################################################################
def test_file_list_generator(target_dir,
                             dir_name="test",
                             prefix_normal="normal",
                             prefix_anomaly="anomaly",
                             ext="wav"):


    normal_files = sorted(
        glob.glob("{dir}/{dir_name}/{prefix_normal}*.{ext}".format(dir=target_dir,
                                                                                dir_name=dir_name,
                                                                                prefix_normal=prefix_normal,
                                                                                ext=ext)))
    normal_labels = numpy.zeros(len(normal_files))
    anomaly_files = sorted(
        glob.glob("{dir}/{dir_name}/{prefix_anomaly}*.{ext}".format(dir=target_dir,
                                                                                dir_name=dir_name,
                                                                                prefix_anomaly=prefix_anomaly,
                                                                                ext=ext)))
    anomaly_labels = numpy.ones(len(anomaly_files))
    files = numpy.concatenate((normal_files, anomaly_files), axis=0)
    labels = numpy.concatenate((normal_labels, anomaly_labels), axis=0)
    print("Test file  num : {num}".format(num=len(files)))

    return files, labels

########################################################################
# Read audio (wav) file 
########################################################################

def file_load(wav_name, mono=False):
    return librosa.load(wav_name, sr=None, mono=mono)


########################################################################
# Extract features from audio waves signal
########################################################################

def file_to_vector_array(file_name,
                         n_mels=64,
                         frames=5,
                         n_fft=1024,
                         hop_length=512,
                         hop_frames = 1,
                         power=2.0):

    # 01 calculate the number of dimensions
    dims = n_mels * frames

    # 02 generate melspectrogram using librosa
    y, sr = file_load(file_name)
    mel_spectrogram = librosa.feature.melspectrogram(y=y,
                                                     sr=sr,
                                                     n_fft=n_fft,
                                                     hop_length=hop_length,
                                                     n_mels=n_mels,
                                                     power=power)

    # 03 convert melspectrogram to log mel energy
    log_mel_spectrogram = 20.0 / power * numpy.log10(mel_spectrogram + sys.float_info.epsilon)

    # 04 calculate total vector size
    vector_array_size = len(log_mel_spectrogram[0, :]) - frames + 1
 
    # 05 skip too short clips
    if vector_array_size < 1:
        return numpy.empty((0, dims))

    # 06 generate feature vectors by concatenating multiframes
    vector_array = numpy.zeros((vector_array_size, dims))
    for t in range(frames):
        vector_array[:, n_mels * t: n_mels * (t + 1)] = log_mel_spectrogram[:, t: t + vector_array_size].T
    
    vector_array  = vector_array[::hop_frames, :]
    return vector_array



########################################################################
# Extract features from list of audio waves signal
########################################################################

def list_to_vector_array(file_list,
                         n_mels=64,
                         frames=5,
                         n_fft=1024,
                         hop_length=512,
                         power=2.0):
    # calculate the number of dimensions
    dims = n_mels * frames

    # iterate file_to_vector_array()
    for idx in tqdm(range(len(file_list))):
        vector_array = file_to_vector_array(file_list[idx],
                                                n_mels=n_mels,
                                                frames=frames,
                                                n_fft=n_fft,
                                                hop_length=hop_length,
                                                power=power)
        if idx == 0:
            dataset = numpy.zeros((vector_array.shape[0] * len(file_list), dims), float)
        dataset[vector_array.shape[0] * idx: vector_array.shape[0] * (idx + 1), :] = vector_array

    return dataset