#!/usr/bin/python3
"""
Tadam !

Tool to preprocess audio in real time from folder coming from recorder SD card 
into time expansion audio splitted in 5s-long chunks to be sent to Vigie-Chiro.
"""

__author__ = "Samuel ORTION"
__copyright__ = "Copyright 2021, Chiro-Canto organisation"
__credits__ = ["Samuel ORTION"]

__license__ = "GPL"
__version__ = "1.0.0"
__maintainer__ = "Samuel ORTION"
__email__ = "samuel.ortion@orange.fr"
__status__ = "Production"
__date__ = "2021-05-17"

import soundfile as sf
import os
import sys
import numpy as np
import getopt
from datetime import datetime
from datetime import timedelta

verbose = False

def usage():
    usage_str = """
        
        Tadam !

        Tool to preprocess audio in real time from folder coming from recorder SD card 
        into time expansion audio splitted in 5s-long chunks to be sent to Vigie-Chiro.        


        Author: Samuel ORTION
        License: GNU GPL v3 or later

        Requirements:
            pysoundfile, numpy
        
        Usage:
            ./tadam.py -i INPUT_FOLDER [ -o OUTPUT_FOLDER ]

        Available options:
            -v (--verbose) - Verbose mode
            -h (--help) - Display usage
            -i (--input) - Specify input folder
            -o (--output) - Specify output folder
            -l (--length) - Specify the max length (in seconds) of output wav
            -z (--zip) - Specify wether the output folder has to be compressed
            -p (--prefix) - Specify the prefix that will be added to the output file name
            -r (--ratio) - Specify samplerate convertion rate

        Example:
            tadam.py -v -i raw -o exp -l 5 -z -p "Car910130-2021-Pass1-Z1-AudioMoth_" -r 0.1 
    """
    print(usage_str)


def rate(infile, outfile, rate_ratio):
    try:
        data, samplerate = sf.read(infile)
        new_samplerate = int(samplerate * rate_ratio)
    except sf.LibsndfileError as e:
        print("ERROR : {0}".format(e))
    try:
        sf.write(outfile, data, new_samplerate)
    except:
        print(f"ERROR while converting sample rate {infile} to {outfile}")

def convert_folder(infolder, outfolder, rate_ratio, prefix):
    try:
        os.mkdir(outfolder)
    except:
        print(f"output folder {outfolder} exists.")
    infiles = os.listdir(infolder)
    n = len(infiles)
    i = 1
    for infile in infiles:
        extension = infile.split(".")[-1]
        if not extension in ["wav", "WAV"]:
            continue
        outfile = prefix + infile.replace('.WAV', '.wav')
        if verbose:
            print(
                f"{str(i).zfill(len(str(n)))}/{n}     Processing {infile} → {outfile}...")
        rate(os.path.join(infolder, infile),
             os.path.join(outfolder, outfile), rate_ratio)
        i += 1


def chunk(outfolder, length):
    for file in os.listdir(outfolder):
        infile = os.path.join(outfolder, file)
        extension = infile.split(".")[-1]
        time = infile.split("_")[-1].split(".")[0]
        if not extension in ["wav", "WAV"]:
            continue
        try:
            data, samplerate = sf.read(infile)
        except:
            print("ERROR: cannot read file " + infile + " for chunking.")
        else:
            if verbose:
                print(f"\t Splitting {infile} in {length}s-long chunks.")
            sections = int(np.ceil(len(data) / samplerate / length))
            for i in range(sections):
                if verbose:
                    print(
                        f"\t {str(i + 1).zfill(len(str(sections)))}/{sections} \t {infile}")

                temp = data[i*samplerate*length: i *
                            samplerate*length+samplerate*length]

                time_file = datetime.strptime(time, "%H%M%S") + timedelta(seconds = 5 * i)
                time_file = str(time_file.time().__format__("%H%M%S"))
                filename = infile.replace(f'{time}.wav', f'{time_file}_000.wav')
                try:
                    sf.write(filename, temp, samplerate)
                except:
                    print(f"ERROR: cannot write chunk file {filename}.")
        os.remove(infile)

def zip_folder(outfolder):
    if verbose:
        print("Compressing .zip archive...")
    os.system('7z -v700m a output.zip exp')
    if verbose:
        print("Archive compressed.")

def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hi:o:vr:l:zp:", [
                                   "help", "input=", "output=", "verbose", "ratio=", "length=", "zip", "prefix="])
    except getopt.GetoptError as err:
        print(err)

    global verbose
    verbose = False
    rate_ratio = 0.1  # Default is x10 time expansion
    infolder = "raw"
    outfolder = "exp"
    compress = False
    prefix = ""
    chunk_length = -1

    for o, a in opts:
        if o in ("-v", "--verbose"):
            verbose = True
        elif o in ("-h", "--help"):
            usage()
            sys.exit()
        elif o in ("-i", "--input"):
            infolder = a
        elif o in ("-o", "--output"):
            outfolder = a
        elif o in ("-r", "--ratio"):
            rate_ratio = float(a)
        elif o in ("-l", "--length"):
            chunk_length = int(a)
        elif o in ("-z", "--zip"):
            compress = True
        elif o in ("-p", "--prefix"):
            prefix = a
        else:
            assert False, "unhandled option"
    convert_folder(infolder, outfolder, rate_ratio, prefix)
    if chunk_length != -1:
        chunk(outfolder, chunk_length)
    if compress:
        zip_folder(outfolder)
    if verbose:
        print("Done.")

if __name__ == "__main__":
    main()