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

__all__ = ['mask2category', 'accuracy', 'lrap', 'kfold_dataframes', 'OneHot', 'MixUp', 'LabelSED']

# Cell
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Beta
from sklearn.metrics import label_ranking_average_precision_score
from sklearn.model_selection import KFold
from IPython.core.debugger import set_trace
from fastai.vision.all import *
from .core import *
from ..vision.losses import *

# Cell
def mask2category(x, y):
    y = TensorAudioLabel(np.nanmax(y.cpu(), axis=(1,2))).type(x.type()).to(x.device)
    return x.float(), y.long()

def accuracy(x, y, before=mask2category, after=lambda o:o):
    f = lambda x,y : (x.argmax(-1) == y).float().mean()
    return after(f(*before(x,y)))

def lrap(x, y, before=mask2category, after=lambda o:o):
    x, y = before(x, y)
    y = F.one_hot(y, x.shape[1]).cpu().numpy()
    x = x.softmax(1).cpu().numpy()
    return label_ranking_average_precision_score(y, x)


# Cell
def kfold_dataframes(df, fold_number, n_splits=5, shuffle=True, random_state=2021):
    kf = KFold(n_splits=n_splits, shuffle=shuffle, random_state=random_state)
    train_idx, valid_idx = list(kf.split(df.index))[fold_number]
    return df.loc[train_idx].reset_index(drop=True), df.loc[valid_idx].reset_index(drop=True)

# Cell
class OneHot(Callback):
    run_valid = False

    def before_batch(self):
        self.learn.yb = tuple([F.one_hot(self.yb[0], self.dls.c).float()])

class MixUp(Callback):
    run_valid = False
    def __init__(self, alpha=0.4, onehot=False):
        self.alpha = alpha
        self.distrib = Beta(alpha, alpha)
        self.onehot = onehot

    def before_batch(self):
        bs      = self.xb[0].shape[0]
        device  = self.xb[0].device
        lambd = self.distrib.sample((self.y.size(0),)).squeeze().to(self.x.device)
        lambd = torch.stack([lambd, 1-lambd], 1).max(1)[0]
        shuffle = torch.randperm(bs).to(device)
        xb1, yb1 = self.xb[0][shuffle], self.yb[0][shuffle]
        a = tensor(lambd).float().view(-1, 1, 1, 1).to(device)
        self.learn.xb = tuple([a*self.xb[0] + (1-a)*xb1])
        a = a.view(-1)
        if self.onehot:
            while len(a.shape) < len(yb1.shape):
                a = a[...,None]
            self.learn.yb = tuple([a*self.learn.yb[0] + (1-a)*yb1])
        else:
            self.learn.yb = tuple([{'yb': self.learn.yb[0], 'yb1': yb1, 'a': a}])

class LabelSED(Callback):
    run_valid = True
    def __init__(self, model_n_rescales):
        self.rescale = 2**model_n_rescales

    def before_batch(self):
        y = self.learn.yb[0]
        y = y[...,::y.shape[-1]//self.rescale].max(2).values.float()
        self.learn.yb = tuple([y])