# AUTOGENERATED! DO NOT EDIT! File to edit: ../00_core.ipynb.

# %% auto 0
__all__ = ['pdlbd_exts', 'get_device', 'is_tool', 'normalize_audio', 'load_audio', 'get_dbmax', 'audio_float_to_int',
           'is_silence', 'batch_it_crazy', 'makedir', 'fast_scandir', 'get_audio_filenames', 'untuple',
           'get_latest_ckpt_old', 'get_latest_ckpt', 'rnd_string', 'get_run_info']

# %% ../00_core.ipynb 4
import torch
import torchaudio
from torchaudio import transforms as T
from torch.nn import functional as F
import numpy as np
from librosa import load as lr_load
from pedalboard.io import AudioFile, get_supported_read_formats
import os
import math
from einops import rearrange
import random
import string
import glob
from pathlib import Path
import warnings

# %% ../00_core.ipynb 5
def get_device(gpu_str=''):
    "utility to suggest which pytorch device to use"
    #return torch.device('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
    device_str = 'cpu'
    if torch.cuda.is_available():
        device_str = 'cuda' if gpu_str=='' else f'cuda:{gpu_str}'
    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): # must check for mps attr if using older pytorch
        device_str = 'mps'
        os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
    return torch.device(device_str)

# %% ../00_core.ipynb 7
def is_tool(name):
    """Check whether `name` is on PATH and marked as executable."""

    # from whichcraft import which
    from shutil import which

    return which(name) is not None

# %% ../00_core.ipynb 9
def normalize_audio(
    audio_in,       # input array/tensor  (numpy or Pytorch)
    norm='global', # global (use max-abs of whole clip) | channel (per-channel norm'd individually) | ''/None
    ):
    "normalize audio, based on the max of the absolute value"
    audio_out = audio_in.clone() if torch.is_tensor(audio_in) else audio_in.copy()  # rudimentary PyTorch/NumPy support
    if ('global' == norm) or  len(audio_in.shape)< 2:
        absmax = abs(audio_in).max()
        audio_out = 0.99*audio_in/absmax if absmax != 0 else audio_in  # 0.99 = just below clipping
    elif 'channel' == norm:
        for c in range(audio_in.shape[0]):  # this loop is slow but sure. TODO: do it fast but still avoid div by zero
            absmax = abs(audio_in[c]).max()
            audio_out[c] = 0.99*audio_in[c]/absmax if absmax != 0 else audio_in[c]  # 0.99 = just below clipping
      #anything else, pass unchanged
    return audio_out

# %% ../00_core.ipynb 19
pdlbd_exts = None      # stores supported pedalboard file extensions. Global so it updates once per run

def load_audio(
    filename:str,     # name of file to load
    sr=48000,         # sample rate in Hz
    verbose=True,     # whether or not to print notices of resampling
    norm='',          # passedto normalize_audio(), see above
    )->torch.tensor:
    "loads an audio file as a torch tensor"
    global pdlbd_exts
    
    if '.mp3' in filename.lower():  # don't rely on torchaudio for mp3s
        pdlbd_exts = get_supported_read_formats() if pdlbd_exts==None else pdlbd_exts
        if '.mp3' in pdlbd_exts:   # first try pedalboard's mp3 support
            with AudioFile(filename) as f:
                audio, in_sr = f.read(f.frames), f.samplerate
        else:
            if verbose: print("Warning: pedalboard mp3 support failed, falling back to librosa")
            audio, in_sr = lr_load(filename, mono=False, sr=sr)
        audio = torch.tensor(audio)
    else:
        audio, in_sr = torchaudio.load(filename)
    if in_sr != sr:
        if verbose: print(f"Resampling {filename} from {in_sr} Hz to {sr} Hz",flush=True)
        resample_tf = T.Resample(in_sr, sr)
        audio = resample_tf(audio)
        
    if norm in ['global','channel']: audio = normalize_audio(audio, norm=norm)
    return audio

# %% ../00_core.ipynb 29
def get_dbmax(
    audio,       # torch tensor of (multichannel) audio
    ):
    "finds the loudest value in the entire clip and puts that into dB (full scale)"
    return 20*torch.log10(torch.flatten(audio.abs()).max()).cpu().numpy()

# %% ../00_core.ipynb 32
def audio_float_to_int(waveform):
    "converts torch float to numpy int16 (for playback in notebooks)"
    return np.clip( waveform.cpu().numpy()*32768 , -32768, 32768).astype('int16')

# %% ../00_core.ipynb 34
def is_silence(
    audio,       # torch tensor of (multichannel) audio
    thresh=-60,  # threshold in dB below which we declare to be silence
    ):
    "checks if entire clip is 'silence' below some dB threshold"
    dBmax = get_dbmax(audio)
    return dBmax < thresh

# %% ../00_core.ipynb 38
def batch_it_crazy(
    x,        # a time series as a PyTorch tensor, e.g. stereo or mono audio
    win_len,  # length of each "window", i.e. length of each element in new batch
    ):
    "(pun intended) Chop up long sequence into a batch of win_len windows"
    if len(x.shape) < 2: x = x.unsqueeze(0)  # guard against 1-d arrays
    x_len = x.shape[-1]
    n_windows = (x_len // win_len) + 1
    pad_amt = win_len * n_windows - x_len  # pad end w. zeros to make lengths even when split
    xpad = F.pad(x, (0, pad_amt))
    return rearrange(xpad, 'd (b n) -> b d n', n=win_len)

# %% ../00_core.ipynb 45
def makedir(
    path:str,              # directory or nested set of directories
    ):
    "creates directories where they don't exist"
    if os.path.isdir(path): return  # don't make it if it already exists
    #print(f"  Making directory {path}")
    try:
        os.makedirs(path)  # recursively make all dirs named in path
    except:                # don't really care about errors
        pass

# %% ../00_core.ipynb 47
def fast_scandir(
    dir:str,  # top-level directory at which to begin scanning
    ext:list  # list of allowed file extensions
    ):
    "very fast `glob` alternative. from https://stackoverflow.com/a/59803793/4259243"
    subfolders, files = [], []
    ext = ['.'+x if x[0]!='.' else x for x in ext]  # add starting period to extensions if needed
    try: # hope to avoid 'permission denied' by this try
        for f in os.scandir(dir):
            try: # 'hope to avoid too many levels of symbolic links' error
                if f.is_dir():
                    subfolders.append(f.path)
                elif f.is_file():
                    if os.path.splitext(f.name)[1].lower() in ext:
                        files.append(f.path)
            except:
                pass 
    except:
        pass

    for dir in list(subfolders):
        sf, f = fast_scandir(dir, ext)
        subfolders.extend(sf)
        files.extend(f)
    return subfolders, files

# %% ../00_core.ipynb 51
def get_audio_filenames(
    paths:list   # directories in which to search
    ):
    "recursively get a list of audio filenames"
    filenames = []
    if type(paths) is str: paths = [paths]
    for path in paths:               # get a list of relevant filenames
        subfolders, files = fast_scandir(path, ['.wav','.flac','.ogg','.aiff','.aif','.mp3'])
        filenames.extend(files)
    return filenames

# %% ../00_core.ipynb 54
def untuple(x, verbose=False):
    """Recursive.  For when you're sick of tuples and lists: 
    keeps peeling off elements until we get a non-tuple or non-list, 
    i.e., returns the 'first' data element we can 'actually use'"""
    if isinstance(x, tuple) or isinstance(x, list): 
        if verbose: print("yea: x = ",x)
        return untuple(x[0], verbose=verbose)
    else:
        if verbose: print("no: x = ",x)
        return x

# %% ../00_core.ipynb 57
def get_latest_ckpt_old(
    dir_tree,            # name of the group of runs without run name or unique identifer
    run_name_prefix='',  # run name without unique identifier for particular run
    verbose=True  # whether to pring message(s)
    ):
    "This will grab the most recent checkpoint filename in dir tree given by name"
    search_string = f"{dir_tree}/{run_name_prefix}*/checkpoints/*.ckpt"
    if verbose: print(f"Looking for latest checkpoint in {search_string}")
    list_of_files = glob.glob(search_string, recursive=True) 
    if [] == list_of_files:
        print("WARNING: No matching checkpoint files found. Starting run from scratch.") 
        return ""
    else:
        return max(list_of_files, key=os.path.getctime)
    
def get_latest_ckpt(
    dir_tree,            # name of the run without unique identifer
    run_name_prefix='',  # unique identifier for particular run
    sim_ckpts=[''],       # string or list of strings. other paths to check under if nothing's in dir_tree
    verbose=True,        # whether to pring message(s)
    ):
    "This will grab the most recent checkpoint filename in dir tree given by name"
    list_of_files = list(Path(dir_tree).glob(f'**/{run_name_prefix}*/checkpoints/*.ckpt'))
    if [] != list_of_files: return max(list_of_files, key=os.path.getctime)
    print(f"   Nothing relevant found in {dir_tree}. Checking also in {sim_ckpts}.")
    if isinstance(sim_ckpts, str): sim_ckpts = [sim_ckpts]
    for pattern in sim_ckpts:
        if verbose: print("   pattern = ",pattern) 
        directories = [dir_path for dir_path in glob.glob(pattern) if not os.path.isfile(dir_path)]
        if verbose: print("   Also checking in ",directories) 
        for directory in directories:
            if verbose: print("     directory = ",directory)
            list_of_files += list(Path(directory).glob(f'**/{run_name_prefix}*/checkpoints/*.ckpt'))
    if [] != list_of_files: return max(list_of_files, key=os.path.getctime)
    warnings.warn("   No matching checkpoint files found anywhere. Starting run from scratch.") 
    return ""

# %% ../00_core.ipynb 60
def rnd_string(n=8): 
    "random letters and numbers of given length. case sensitive"
    raise DeprecationWarning("Better to generate random string in SLURM script")
    return ''.join(random.choice(string.ascii_letters+string.digits) for i in range(n))

def get_run_info(run_name, verbose=True):
    """
    parses run_name into (ideally) prefix & id using underscore as separator, and/or fills in missing info if needed
    NOTE: do not trust generated strings to be same on other processes
    """
    run_info = run_name.split('_')
    prefix = run_info[0]
    if len(run_info)>1: 
        run_id = run_info[-1] 
    else:
        run_id = rnd_string()
        if verbose: print(f"WARNING: generating random run_id as {run_id}. Might be different on different process(es).")
        raise DeprecationWarning("Instead, generate random string in SLURM script")
    new_run_name = f"{prefix}_{run_id}" if prefix !='' else f"{run_id}" 
    return {'prefix': prefix, 'id':run_id, 'run_name':new_run_name}
