# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01audio_core.ipynb (unless otherwise specified).

__all__ = ['SAMPLE_FILE', 'TensorAudio', 'TensorAudioLabel', 'load_npy', 'sample_file', 'melspectrogram', 'show_sample',
           'audio2npy']

# Cell
import numpy as np
import torch
import matplotlib.pyplot as plt
from fastai.torch_core import TensorBase
from nnAudio import Spectrogram
from fastcore.all import *
from fastai.vision.all import *
import librosa

# Cell
SAMPLE_FILE = '/kaggle/kaggle_rainforest_audio/data/npy32000/train/000316da7.npy'

# Cell
class TensorAudio(TensorBase): pass
class TensorAudioLabel(TensorBase): pass

def load_npy(fn):
    return TensorAudio(np.load(fn))

def sample_file():
    return SAMPLE_FILE

# Cell
def melspectrogram(wav, sample_rate=22050, n_fft=2048, n_mels=128, hop_length=512,
                   window='hann', center=True, pad_mode='reflect', power=2.0, htk=False,
                   fmin=0.0, fmax=None, norm=1, trainable_mel=False, trainable_STFT=False,
                   verbose=False, eps=1e-6, cuda=False, log=True, **kwargs):
    s =  Spectrogram.MelSpectrogram(sr=sample_rate, n_fft=n_fft, n_mels=n_mels,
        hop_length=hop_length, window=window, center=center, pad_mode=pad_mode, power=power,
        htk=htk, fmin=fmin, fmax=fmax, norm=norm, trainable_mel=trainable_mel,
        trainable_STFT=trainable_STFT, verbose=verbose, **kwargs)
    if cuda: s = s.cuda()
    return torch.log(s(wav) + eps) if log else s(wav)

def show_sample(file=sample_file(), tfms=lambda x : x, sample_rate=48_000, hop_length=512, width=512,
                figsize=(10,4), vmin=-10, vmax=2):
    plt.figure(figsize=figsize)
    plt.imshow(melspectrogram(tfms(load_npy(file)[:width*hop_length]), sample_rate).squeeze(),
               cmap='RdYlGn_r', vmin=vmin, vmax=vmax)
    plt.colorbar(orientation='horizontal')
    plt.show()

# Cell
def audio2npy(file, path_save, sample_rate=32_000):
    path_save.mkdir(exist_ok=True, parents=True)
    path_save.mkdir(exist_ok=True, parents=True)
    wave, _ = librosa.load(file, sr=sample_rate)
    np.save(path_save/f'{file.stem}.npy', wave)