import glob
import soundfile as sf
import scipy.signal as sg
import argparse
import sys
import os
import numpy as np

from tqdm.contrib.concurrent import process_map
from multiprocessing import cpu_count

TARGET_SR = 256000

PATH_DATA = "/nfs/NAS3/SABIOD/SITE/ANTILLES_CCS_2021/"
PATH_OUT = "/nfs/NAS3/SABIOD/SITE/ANTILLES_CCS_2021_RESAMPLE_256kHz/"

def get(file):
    try:
      ob = sf.SoundFile(file)
    except:
      return file

    if ob.samplerate != 256000:
        try:
            sig = ob.read()
        except:
            return file

        if sig.ndim > 1:
          sig = sig[:, 0]

        if len(sig) == 0:
          return file

        out_file = file.replace(PATH_DATA, PATH_OUT)
        if os.path.exists(out_file):
            return ""

        resample_sig = sg.resample(sig, int((len(sig) * TARGET_SR) / ob.samplerate))
        os.makedirs(os.path.dirname(out_file), exist_ok=True)
        sf.write(out_file, resample_sig, TARGET_SR, ob.subtype)

    return ""



parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                     description="""TODO""")
parser.add_argument("name_session", type=str, help="folder contening audiofiles which must be resampled (example : 'LOT2/ANGUILLA_20210128_20210318')")
parser.add_argument('--number_cpu', type=int, default=8) # number of  cpu used
args = parser.parse_args()

name_session = args.name_session

l_files = glob.glob(PATH_DATA + name_session + '/*.wav')

print("Session :" + name_session)
nb_cpu = np.min([cpu_count(), args.number_cpu])
l_invalid_files = process_map(get, l_files, max_workers=nb_cpu, chunksize=1)
print("All files can't be opened :" + str(list(filter(None, l_invalid_files))))