#!/usr/bin/env python3
import argparse, os, shutil, subprocess
from multiprocessing import Pool, cpu_count

#os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
#import tensorflow as tf
#tf.get_logger().setLevel('FATAL')

from essentia.standard import *
essentia.log.infoActive = False

DIRS = ['_new', '_new-dubstep']
RATE = 44100 

parser = argparse.ArgumentParser()
parser.add_argument('dirs', nargs='*', default=DIRS, help='the list of directories')
args = parser.parse_args()

def calculate_key(audio_path):
    import essentia.streaming as ess
    import essentia

    loader = ess.MonoLoader(filename=audio_path)
    framecutter = ess.FrameCutter(frameSize=4096, hopSize=2048, silentFrames='noise')
    windowing = ess.Windowing(type='blackmanharris62')
    spectrum = ess.Spectrum()
    spectralpeaks = ess.SpectralPeaks(orderBy='magnitude',
                                      magnitudeThreshold=0.00001,
                                      minFrequency=20,
                                      maxFrequency=3500,
                                      maxPeaks=60)
    # Use default HPCP parameters for plots.
    # However we will need higher resolution and custom parameters for better Key estimation.
    hpcp = ess.HPCP()
    hpcp_key = ess.HPCP(size=36, # We will need higher resolution for Key estimation.
                        referenceFrequency=440, # Assume tuning frequency is 44100.
                        bandPreset=False,
                        minFrequency=20,
                        maxFrequency=3500,
                        weightType='cosine',
                        nonLinear=False,
                        windowSize=1.)
    
    key = ess.Key(profileType='edmm', # Use profile for electronic music.
                  numHarmonics=4,
                  pcpSize=36,
                  slope=0.6,
                  usePolyphony=True,
                  useThreeChords=True)
    
    # Use pool to store data.
    pool = essentia.Pool()
    
    # Connect streaming algorithms.
    loader.audio >> framecutter.signal
    framecutter.frame >> windowing.frame >> spectrum.frame
    spectrum.spectrum >> spectralpeaks.spectrum
    spectralpeaks.magnitudes >> hpcp.magnitudes
    spectralpeaks.frequencies >> hpcp.frequencies
    spectralpeaks.magnitudes >> hpcp_key.magnitudes
    spectralpeaks.frequencies >> hpcp_key.frequencies
    hpcp_key.hpcp >> key.pcp
    hpcp.hpcp >> (pool, 'tonal.hpcp')
    key.key >> (pool, 'tonal.key_key')
    key.scale >> (pool, 'tonal.key_scale')
    key.strength >> (pool, 'tonal.key_strength')
    
    # Run streaming network.
    essentia.run(loader)
    return pool['tonal.key_key'], pool['tonal.key_scale']

def calculate_bpm(audio_path):
    audio44100 = MonoLoader(filename=audio_path, sampleRate=44100, resampleQuality=4)()
    start44100 = int(len(audio44100)*0.2)
    end44100 = int(len(audio44100)*0.8)
    
    rhythm_extractor = RhythmExtractor2013(method="multifeature")
    bpm, _, _, _, _ = rhythm_extractor(audio44100[start44100:end44100])
    if bpm < 100:
        bpm = bpm * 2
    return bpm

def reencode(audio_path):
    if '.wav' in audio_path[-4:]:
        subprocess.run("ffmpeg -hide_banner -loglevel error -y -i \"" + audio_path + "\" -map 0:a -c:a pcm_s16le -ar " + str(RATE) + " -write_xing 0 -id3v2_version 0 -map_metadata -1 -bitexact \"" + audio_path + ".reenc.wav\"", shell=True)
        shutil.move(audio_path + ".reenc.wav", audio_path)
        return audio_path
    elif '.flac' in audio_path[-5:]:
        subprocess.run("ffmpeg -hide_banner -loglevel error -y -i \"" + audio_path + "\" -map 0:a -c:a pcm_s16le -ar " + str(RATE) + " -write_xing 0 -id3v2_version 0 -map_metadata -1 -bitexact \"" + audio_path.replace('.flac', '.wav') + "\"", shell=True)
        os.remove(audio_path)
        return audio_path.replace('.flac', '.wav')
    elif '.mp3' in audio_path[-4:]:
        subprocess.run("ffmpeg -hide_banner -loglevel error -y -i \"" + audio_path + "\" -map 0:a -c:a copy -write_xing 0 -id3v2_version 0 -map_metadata -1 \"" + audio_path + ".reenc.mp3\"", shell=True)
        shutil.move(audio_path + ".reenc.mp3", audio_path)
        return(audio_path)
    else:
        print("Can't reencode", audio_path)
        return None
 
def rmtags(audio_path):
    if '.wav' in audio_path[-4:]:
        subprocess.run("id3v2 -D \"" + audio_path + "\" > /dev/null", shell=True)
    elif '.mp3' in audio_path[-4:]:
        subprocess.run("id3v2 -D \"" + audio_path + "\" > /dev/null", shell=True)
        subprocess.run("apetag -m erase -i \"" + audio_path + "\" > /dev/null", shell=True)
    else:
        print("Can't rmtags'", audio_path)

def process(audio_path):
    #key, scale = calculate_key(audio_path)
    audio_path = reencode(audio_path)
    if audio_path is not None:
        bpm = "{:.2f}".format(calculate_bpm(audio_path))
        if bpm in audio_path:
            name = audio_path
        else:
            rmtags(audio_path)
            name = audio_path[:-4] + " - " + bpm + audio_path[-4:]
            shutil.move(audio_path, name)
        return name

files = []

for dir in args.dirs:
    for file in os.listdir(dir):
        files.append(dir + "/" + file)

NICENESS=15

with Pool(initializer=os.nice, initargs=(NICENESS,)) as pool:
    for result in pool.imap_unordered(process, files):
        print(result)

print('done\n')