# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/audio.04_learner.ipynb.

# %% ../../nbs/audio.04_learner.ipynb 1
# Contains code from https://github.com/fastaudio/fastaudio:
# fastaudio - MIT License - Copyright (c) 2020 scart97

# %% ../../nbs/audio.04_learner.ipynb 3
from __future__ import annotations

from pathlib import Path

from fastcore.dispatch import retain_type

from fastai.callback.core import Callback
from fastai.callback.fp16 import MixedPrecision
from fastai.data.core import DataLoaders
from fastai.learner import Metric, Learner, defaults
from fastai.optimizer import Optimizer, OptimWrapper

from .core import TensorSpec, TensorMelSpec
from .data import MelSpectrogram, Spectrogram
from ..optimizer.fused import Adam
from ..imports import *

# %% auto 0
__all__ = ['StackSpecCallback', 'audio_learner']

# %% ../../nbs/audio.04_learner.ipynb 5
class StackSpecCallback(Callback):
    "Stacks tuples of TensorSpec or TensorMelSpec. ToDo: add resizing"
    order = MixedPrecision.order-1
    def before_batch(self):
        xb = L(self.xb)
        idx = xb.argwhere(lambda x: isinstance(x, (TensorSpec, TensorMelSpec)))
        ts = []
        for i in idx:
            ts.append(xb[i])
        stacked = torch.stack(ts, dim=2)
        xb = retain_type(torch.flatten(stacked, 1, 2), xb[i])
        self.learn.xb = tuple(xb)

# %% ../../nbs/audio.04_learner.ipynb 7
# TODO opt_func type hint should be Callable[..., Optimizer|OptimWrapper]|OptimWrapper
# but nbdev 2.3.9 currently doesn't support this type hint

def audio_learner(
    dls:DataLoaders, # `DataLoaders` containing fastai or PyTorch `DataLoader`s
    model:nn.Module, # PyTorch model for training or inference
    loss_func:nn.Module|Callable[..., Tensor]|None=None, # Loss function. Defaults to `dls` loss
    opt_func:Optimizer|OptimWrapper=Adam, # Optimization function for training
    lr:float|slice=defaults.lr, # Default learning rate
    splitter:Callable[[nn.Module], list[Tensor]]=trainable_params, # Split model into parameter groups. Defaults to one parameter group
    cbs:Listified[Callback]|None=None, # `Callback`s to add to `Learner`
    metrics:Listified[Metric]|None=None, # `Metric`s to calculate on validation set
    path:str|Path|None=None, # Parent directory to save, load, and export models. Defaults to `dls` `path`
    model_dir:str|Path='models', # Subdirectory to save and load models
    wd:float|None=None, # Default weight decay
    wd_bn_bias:bool=False, # Apply weight decay to normalization and bias parameters
    train_bn:bool=True, # Train frozen normalization layers
    moms:tuple[float,float,float]=(0.95,0.85,0.95), # Default momentum for schedulers
    default_cbs:bool=True # Include default `Callback`s
) -> Learner:
    "An Audio specific Learner that stacks tuples of `TensorSpec` or `TensorMelSpec`"
    detuple = False
    for i in range(len(dls.train.after_batch.fs)):
        if not detuple and isinstance(dls.train.after_batch[i], (Spectrogram, MelSpectrogram)):
            detuple = is_listy(dls.train.after_batch[i].n_fft)

    if detuple:
        if cbs is None:
            cbs = StackSpecCallback()
        else:
            cbs = L(cbs) + L(StackSpecCallback())

    return Learner(dls=dls, model=model, loss_func=loss_func, opt_func=opt_func, lr=lr, splitter=splitter, cbs=cbs,
                    metrics=metrics, path=path, model_dir=model_dir, wd=wd, wd_bn_bias=wd_bn_bias, train_bn=train_bn,
                    moms=moms, default_cbs=default_cbs)
