# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/kaggle_rfcx-species-audio-detection.ipynb (unless otherwise specified).

__all__ = ['audio_augment', 'train', 'get_preds', 'test', 'main']

# Cell
import numpy as np
import torch
import torch.nn as nn
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime
from fastprogress import progress_bar
from IPython.core.debugger import set_trace
import gc
from fastscript import *
from fastcore.all import *
from fastai.vision.all import *
from ..core import *
from ..audio.core import *
from ..audio.augmentations import *
from ..audio.dataset import *
from ..vision.models import *
from ..vision.losses import *
from ..audio.util import *

# Cell
def audio_augment(sample_rate, p=0.25):
    return Pipeline([
        ClippingDistortion(sample_rate, max_percentile_threshold=10, p=p),
        PitchShift(sample_rate, min_semitones=-8, max_semitones=8, p=p),
    ])

# Cell
def train(sample_rate, num_classes, fold, n_epochs, lr, wd, tile_width, bs, aug_ps,
          model_name, loss_func, plot, load_checkpoint=None, lr_find=False, head_ps=0.8,
          mixup=False, n_mels=128, hop_length=512, model='resnest50'):
    seed_everything()
    cbs = []
    path = Path('/kaggle/kaggle_rainforest_audio/data')
    rename_cols = RenameColumns(id='recording_id', label='species_id', tmin='t_min',
                                tmax='t_max',fmin='f_min', fmax='f_max')

    df = Pipeline([load_dataframe, rename_cols, group_labels])(path/'train_tp.csv')

    train_df, valid_df = kfold_dataframes(df, fold)

    tfms = partial(apply_augmentations, augs_pipeline=audio_augment(sample_rate, p=aug_ps))

    train_data = Datasets(items=train_df, tfms=partial(create_dataset_item, path=path,
                                               sample_rate=sample_rate, tile_width=tile_width,
                                               n_mels=n_mels, hop_length=hop_length))

    valid_data = Datasets(items=valid_df, tfms=partial(create_dataset_item, path=path,
                                               sample_rate=sample_rate, tile_width=tile_width,
                                               n_mels=n_mels, hop_length=hop_length))
    train_dl = DataLoader(
        train_data, bs=bs, do_batch=reorganize_batch, shuffle=True,
        num_workers=8, after_item=tfms,
        after_batch=MelSpectrogram(sample_rate,n_mels=n_mels,hop_length=hop_length))

    valid_dl = DataLoader(
        valid_data, bs=bs, do_batch=reorganize_batch, num_workers=8,
        after_batch=MelSpectrogram(sample_rate, n_mels=n_mels,hop_length=hop_length))

    dls = DataLoaders(train_dl, valid_dl)
    dls.device = torch.device("cuda:0")

    if plot:
        xb, yb = dls.one_batch()
        show_augmentations(train_data, train_dl, sample_rate=sample_rate)

    model = get_model(model, num_classes=num_classes, head_ps=head_ps, in_channels=1)

    if mixup:
        cbs.append(MixUp(0.4))
        loss_func += '_mixup'

    def before_loss(x,y):
        x,y=mask2category(x,y)
        return x, y

    def after_loss(loss, y):
        return loss

    loss = get_loss(loss_func, before=before_loss, after=after_loss)
    print('Loss function: ', loss_func)

    learn = Learner(dls, model, loss_func=loss, metrics=[accuracy, lrap], cbs=cbs)
    learn.to_fp16(clip=0.5);

    if load_checkpoint is not None:
        learn.load(path.parent/f'models/{load_checkpoint}_fold{fold}')
        print('Load model ', path.parent/f'models/{load_checkpoint}_fold{fold}')

    if lr_find: learn.lr_find()

    learn.fit_one_cycle(n_epochs, lr, wd=wd, div_final=10, div=10)
    learn.save(path.parent/f'models/{model_name}_fold{fold}')
    print(f'Model saved to', path.parent/f'models/{model_name}_fold{fold}')

# Cell
def get_preds(dataloader, model, device=torch.device("cuda:0"), max_reduce=True):
    model.eval().to(device)
    with torch.no_grad():
        preds, ys = [], []
        for x, y in progress_bar(dataloader):
            if max_reduce:
                pred = model(x).max(0).values[None]
            else:
                pred = model(x)[None]
            preds.append(pred.cpu())
            ys.append(y.cpu())
        preds = torch.cat(preds, dim=0)
        ys = torch.cat(ys, dim=0)
    return preds, ys

def test(sample_rate, num_classes, tile_width, model_name, ens_folds, head_ps=0.8,
         n_mels=128, hop_length=512, save_preds=False, model='resnest50'):
    bs = 1
    _path_save = Path('preds')
    _path_save.mkdir(exist_ok=True)
    max_reduce = not save_preds

    preds_ens, preds_valid, ys_valid = [], [], []
    preds_train, ys_train = [], []
    for fold in ens_folds:
        seed_everything()
        path = Path('/kaggle/kaggle_rainforest_audio/data')
        rename_cols_test = RenameColumns(id='recording_id')
        rename_cols_valid = RenameColumns(id='recording_id', label='species_id', tmin='t_min',
                                    tmax='t_max',fmin='f_min', fmax='f_max')

        df = Pipeline([load_dataframe, rename_cols_valid, group_labels])(path/'train_tp.csv')

        train_df, valid_df = kfold_dataframes(df, fold)

        test_df = Pipeline([load_dataframe, rename_cols_test])(path/'sample_submission.csv')
        fp_df = Pipeline([load_dataframe, rename_cols_valid, group_labels])(path/'train_fp.csv')

        datasets = [Datasets(items=dataframe, tfms=partial(create_dataset_item, path=path,
            sample_rate=sample_rate, tile_width=None, n_mels=n_mels, hop_length=hop_length))
            for dataframe in [train_df, valid_df, test_df, fp_df]]

        dls = DataLoaders(*[DataLoader(dataset, bs=bs, do_batch=reorganize_batch, num_workers=8,
                        after_batch=Pipeline([MelSpectrogram(sample_rate, n_mels=n_mels,
                                    hop_length=hop_length), TilesTransform(tile_width)]))
             for dataset in datasets])

        model = get_model(model, num_classes=num_classes, head_ps=head_ps, in_channels=1,
                          pretrained=False)

        dls.device = torch.device("cuda:0")
        learn = Learner(dls, model, loss_func=cross_entropy, metrics=[accuracy, lrap])
        learn.to_fp16(clip=0.5);
        learn.load(path.parent/f'models/{model_name}_fold{fold}')
        print('Load model: ', path.parent/f'models/{model_name}_fold{fold}')

        preds, ys = get_preds(dls[1], model, max_reduce=max_reduce)
        np.save(_path_save/f'{model_name}_{tile_width}_fold{fold}_valid.npy',
                {'preds':preds, 'ys':ys})
        if not max_reduce: preds = preds.max(1).values
        preds_valid.append(preds)
        ys_valid.append(ys)

        preds, ys = get_preds(dls[2], model, max_reduce=max_reduce)
        np.save(_path_save/f'{model_name}_{tile_width}_fold{fold}_test.npy',
                {'preds':preds, 'ys':ys})
        if not max_reduce: preds = preds.max(1).values
        preds_ens.append(preds[None])

    preds_valid, ys_valid = torch.cat(preds_valid), torch.cat(ys_valid)
    valid_score = lrap(preds_valid, ys_valid.long().squeeze(), before=lambda *o:o)
    print(f'Validation score: {valid_score:.3f}')

    preds_ens = torch.cat(preds_ens).mean(0).softmax(-1)
    test_df = Pipeline([load_dataframe])(path/'sample_submission.csv')

    for i in range(preds_ens.shape[1]):
        test_df.loc[:, f's{i}'] = preds_ens[:,i]

    tstr = datetime.now().strftime('%Y%m%d%H%M')
    test_df.to_csv(path.parent/f'subs/submission_{tstr}.csv',
                   index=False)
    print('Submission file saved: ', path.parent/f'subs/submission_{tstr}.csv')

# Cell
@call_parse
def main(fold:Param('Fold number', int)=0,
         n_epochs:Param('Number of training epochs', int)=30,
         lr:Param('Learning rate', float)=1e-3,
         wd:Param('Weight decay', float)=3e-2,
         loss_func:Param('Loss function', str)='cross_entropy',
         tile_width:Param('Tile width', int)=1024,
         sample_rate:Param('Sample rate', int)=32000,
         n_mels:Param('Spectrogram n_mels', int)=128,
         hop_length:Param('Spectrogram hop_length', int)=640,
         bs:Param('Batch size', int)=32,
         accumulate_gradients:Param('Batch size for gradient accumulation', int)=None,
         aug_ps:Param('Augmentation probability', float)=0.25,
         model:Param('Name of model architecture', str)='densenet121',
         model_name:Param('Name of parameters file', str)='model_n0',
         ens_folds:Param('Folds to use for ensemble', list)=[0,1,2,3,4],
         run_test:Param('Run test prediction (default is train)', str)=False,
         load_checkpoint:Param('Load model checkpoint before new train loop', str)=None,
         head_ps:Param('Model head dropout probability', float)=0.8,
         mixup:Param('Use mixup', str)=True,
         save_preds:Param('Save model predictions for train and validation', str)=False):

    num_classes = 24
    if run_test in [True, 'true', 'True']: run_test = True
    if mixup in [True, 'true', 'True']: mixup = True
    if save_preds in [True, 'true', 'True']: save_preds=True
    if run_test:
        test(sample_rate, num_classes, tile_width, model_name, ens_folds, head_ps=head_ps,
             n_mels=n_mels, hop_length=hop_length, save_preds=save_preds, model=model)
    else:
        train(sample_rate, num_classes, fold, n_epochs, lr, wd, tile_width, bs, aug_ps,
          model_name, loss_func, plot=False, load_checkpoint=load_checkpoint, lr_find=False,
          head_ps=head_ps, mixup=mixup, n_mels=n_mels, hop_length=hop_length, model=model)