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

__all__ = ['MatthewsCorrCoefBinary', 'SQNyr', 'avg_R', 'perc', 'win_perc', 'SQN100', 'tstat', 'tstat_zero',
           'total_est_Rs', 'ΔR_bin_perc', 'avgR_class', 'perc_class', 'win_perc_class', 'tstat_bin_ext',
           't_stat_binary', 'get_task_metrics']

# Cell
from .imports import *
from fastai.metrics import *

# Cell
from scipy.stats import ttest_ind

# Cell
mk_class('ActivationType', **{o:o.lower() for o in ['No', 'Sigmoid', 'Softmax', 'BinarySoftmax']},
         doc="All possible activation classes for `AccumMetric")

# Cell
def MatthewsCorrCoefBinary(thresh=.5, sample_weight=None):
    "Matthews correlation coefficient for single-label classification problems"
    return skm_to_fastai(skm.matthews_corrcoef, activation=ActivationType.BinarySoftmax, thresh=thresh, sample_weight=sample_weight)

# Cell
def _avg_R(inp, targ):
    "Compute average R"
    return targ[inp > 0].mean()
avg_R = AccumMetric(_avg_R, to_np=True)

def _perc(inp, targ):
    "Compute percent"
    return (inp > 0).mean()
perc = AccumMetric(_perc, to_np=True)

def _win_perc(inp, targ):
    "Compute winning percent"
    return (targ[inp > 0] > 0).mean()
win_perc = AccumMetric(_win_perc, to_np=True)

def _SQN100(inp, targ):
    "Compute SQN-100"
    if (inp > 0).sum() >= 2: sqn100 = 10 * targ[inp > 0].mean() / targ[inp > 0].std()
    else: return np.nan
    return sqn100
SQN100 = AccumMetric(_SQN100, to_np=True)

def _SQNyr(inp, targ, per_year=100):
    "Compute SQN-year"
    if (inp > 0).sum() >= 2: return np.sqrt(per_year * (inp > 0).sum() / len(inp)) * targ[inp > 0].mean() / targ[inp > 0].std()
    else: return np.nan

def SQNyr(per_year=100):
    return AccumMetric(_SQNyr, per_year=per_year, to_np=True)

def _tstat(inp, targ):
    "Compute t-stat based on 2 independent sample distributions (regression and binary tasks)"
    if inp.ndim > 1: inp = inp[:, -1]
    return scipy.stats.ttest_ind(targ[inp > 0], targ[inp <= 0])[0]
tstat = AccumMetric(_tstat, to_np=True)

def _tstat_zero(inp, targ):
    "Compute t-stat based on 2 independent sample distributions (regression and binary tasks)"
    if inp.ndim > 1: inp = inp[:, -1]
    return scipy.stats.ttest_ind(targ[inp > 0], torch.zeros((inp > 0).sum()))[0]
tstat_zero = AccumMetric(_tstat, to_np=True)

# Cell
def _total_est_Rs(inp, targ, avg_win=1, avg_loss=-1, per_year=100, thresh=0.5):
    "Compute total Rs / year"
    win_perc = targ[inp].mean() if inp.sum() > 0 else 0
    total = per_year * inp.mean()
    return int(win_perc * total * avg_win + (1 - win_perc) * total * avg_loss)

def total_est_Rs(avg_win=1, avg_loss=-1, per_year=100, thresh=0.5):
    return AccumMetric(_total_est_Rs, flatten=False, avg_win=avg_win, avg_loss=avg_loss, per_year=per_year,
                       to_np=True, activation=ActivationType.BinarySoftmax, thresh=thresh)

# Cell
def _ΔR_bin_perc(inp, targ, value=None, thresh=0.5):
    "Τotal incremental Rs relative to optimal (1) and average (0)"
    ΔR_target = value[(targ == 1)].mean() - value.mean()
    ΔR_pred = value[(inp == 1)].mean() - value.mean()
    return min(1, ΔR_pred/ΔR_target) * min(1, inp.sum()/targ.sum())

def ΔR_bin_perc(value=None, thresh=0.5):
    return AccumMetric(_ΔR_bin_perc, value=value, flatten=False, to_np=True, activation=ActivationType.BinarySoftmax, thresh=thresh)

# Cell
def _avgR_class(inp, targ, value=None):
    "Compute average R in classification tasks - unshuffled"
    return value[inp].mean()

def avgR_class(value, thresh=.5):
    return AccumMetric(_avgR_class, value=value, flatten=False, activation=ActivationType.BinarySoftmax, to_np=True, thresh=thresh)

def _perc_class(inp, targ):
    "Compute percent > .5"
    return inp.mean()

def perc_class(thresh=.5):
    return AccumMetric(_perc_class, flatten=False, activation=ActivationType.BinarySoftmax, to_np=True, thresh=thresh)

def _win_perc_class(inp, targ, value=None):
    "Compute winning percent in classification tasks - unshuffled"
    return (value[inp] > 0).mean()

def win_perc_class(value, thresh=.5):
    return AccumMetric(_win_perc_class, value=value, flatten=False, activation=ActivationType.BinarySoftmax, to_np=True, thresh=thresh)

# Cell
def _tstat_bin_probas(inp, targ):
    "Compute t-stat based on binary probas"
    return scipy.stats.ttest_ind(inp[targ == 1], inp[targ == 0])[0]
t_stat_binary = AccumMetric(_tstat_bin_probas, to_np=True, activation=ActivationType.BinarySoftmax)

def _tstat_bin_ext(inp, targ, ext=None):
    "Compute t-stat for an external variable based on binary probas"
    return scipy.stats.ttest_ind(ext[inp == 1], ext[inp == 0])[0]

def tstat_bin_ext(ext, thresh=.5):
    return AccumMetric(_tstat_bin_ext, ext=ext, flatten=False, activation=ActivationType.BinarySoftmax, to_np=True, thresh=thresh)

# Cell
def get_task_metrics(dls, binary_metrics=None, multi_class_metrics=None, regression_metrics=None, verbose=True):
    if dls.c == 2:
        pv('binary-classification task', verbose)
        return binary_metrics
    elif dls.c > 2:
        pv('multi-class task', verbose)
        return multi_class_metrics
    else:
        pv('regression task', verbose)
        return regression_metrics