# AUTOGENERATED! DO NOT EDIT! File to edit: ../04_spectrofu.ipynb.

# %% auto 0
__all__ = ['save_stft', 'process_one_file', 'main']

# %% ../04_spectrofu.ipynb 6
import argparse 
from glob import glob 
from pathlib import Path
import os 
import math
from multiprocessing import Pool, cpu_count, Barrier
from functools import partial
from tqdm.contrib.concurrent import process_map  
import torch
import torchaudio
from .core import is_silence, load_audio, makedir
from .viz import audio_spectrogram_image

# %% ../04_spectrofu.ipynb 7
def save_stft(
    audio:torch.tensor,  # long audio file to be chunked
    new_filename:str    # stem of new filename(s) to be output as spectrogram images
    ):
    "coverts audio to stft image and saves it"
    im = audio_spectrogram_image(audio, justimage=True)  # should already be a PIL image
    print(f"saving new file = {new_filename}")
    im.save(new_filename)
    return

# %% ../04_spectrofu.ipynb 8
def process_one_file(
    filenames:list,      # list of filenames from which we'll pick one
    args,                # output of argparse
    file_ind             # index from filenames list to read from
    ):
    "this turns one audio file into a spectrogram.  left channel only for now"
    filename = filenames[file_ind]  # this is actually input_path+/+filename
    output_path, input_paths = args.output_path, args.input_paths
    new_filename = None
    
    for ipath in input_paths: # set up the output filename & any folders it needs
        if ipath in filename: # this just avoids repeats/ weirdness.
            last_ipath = ipath.split('/')[-1]           # get the last part of ipath
            clean_filename = filename.replace(ipath,'') # remove all of ipath from the front of filename
            new_filename = f"{output_path}/{last_ipath}/{clean_filename}".replace('//','/')
            new_filename = str(Path(new_filename).with_suffix(".png"))  # give it file extension for image
            makedir(os.path.dirname(new_filename))      # we might need to make a directory for the output file
            break
    
    if new_filename is None:
        print(f"ERROR: Something went wrong with name of input file {filename}. Skipping.",flush=True) 
        return 

    try:
        audio = load_audio(filename, sr=args.sr)
        save_stft(audio, new_filename)
    except Exception as e: 
        print(f"Some kind of error happened with {filename}, either loading or writing images. Skipping.", flush=True)

    return


def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--sr', type=int, default=48000, help='Output sample rate')
    parser.add_argument('--workers', type=int, default=min(32, os.cpu_count() + 4), help='Maximum number of workers to use (default: all)')
    parser.add_argument('output_path', help='Path of output for spectrogram-ified data')
    parser.add_argument('input_paths', nargs='+', help='Path(s) of a file or a folder of files. (recursive)')
    args = parser.parse_args()

    print(f"  output_path = {args.output_path}")

    print("Getting list of input filenames")
    filenames = get_audio_filenames(args.input_paths) 
    n = len(filenames)   
    print(f"  Got {n} input filenames") 

    print("Processing files (in parallel)")
    wrapper = partial(process_one_file, filenames, args)
    r = process_map(wrapper, range(0, n), chunksize=1, max_workers=args.workers)  # different chunksize used by tqdm. max_workers is to avoid annoying other ppl

    print("Finished")
