from features import *
import gc

print("Garbage collection thresholds:",
                    gc.get_threshold())


def extract(args):
    af, output_directory, overwrite = args
    filename = af.split(os.sep)[-1]
    output_file = os.path.splitext(filename)[0]
    output_file = os.path.join(output_directory, output_file)

    if os.path.exists(output_file + ".npz") and not overwrite:
        print('Skipping {}. Already exists.'.format(output_file + ".npz"))
        return

    t_item_tot = datetime.datetime.now()

    print(filename, ' : processing...')

    y, _sr = load(af)
    y = to_mono(y)
    sr = 22050
    y = resample(y, _sr, sr)

    output = dict()
    output['linspec_mag'], output['linspec_phase'] = linspec(y)
    output['melspec'] = melspec(y, sr=sr)
    output['logspec'] = logspec(y, sr=sr)
    output['hcqt_mag'], output['hcqt_phase'] = hcqt(y, sr=sr)
    output['vggish_melspec'] = vggish_melspec(y, sr=sr)

    # high-level
    output['percussive_ratio'], output['percussive_rms'], output['total_rms'] = percussive_ratio(y, margin=3.0)
    output['onset_strength'] = onset_strength(y, detrend=True)
    output['tempogram'] = tempogram(y)
    output['onset_patterns'] = onset_patterns(y, sr=sr)

    np.savez_compressed(output_file, **output)

    print(filename, ' : total processing time :', datetime.datetime.now()-t_item_tot)

    del y
    for obj in output:
        del output[obj]
    del output

    collected = gc.collect()
    print("Garbage collector: collected",
          "%d objects." % collected)
    gc.get_count()

def process_arguments(args):
    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument('csv_file', type=str,
                        help="The path to the CSV file (ID, path).")

    parser.add_argument('output_directory', type=str,
                        help="Path to output directory.")

    parser.add_argument('cpus', type=int,
                        help="Number of processors")

    parser.add_argument('--overwrite', action='store_true', default=False)

    return parser.parse_args(args)

if __name__ == '__main__':
    import soundfile
    import csv

    params = process_arguments(sys.argv[1:])

    csv_file = open(params.csv_file, newline='')
    reader = csv.reader(csv_file, delimiter=';', quotechar='|')
    audio_files = [row[1] for row in reader if row[1].endswith(('.wav', '.mp3'))]

    if not os.path.exists(params.output_directory):
        os.makedirs(params.output_directory)

    t_tot = datetime.datetime.now()

    p = Pool(processes=params.cpus)
    max_ = len(audio_files)
    ret = p.map(extract, [(af, params.output_directory, params.overwrite) for af in audio_files])

    print('total time', datetime.datetime.now()-t_tot)

