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

__all__ = []

# Cell
from .imports import *
from .utils import random_shuffle
from .data.core import *
from .inference import get_X_preds
from fastai.learner import *

# Cell
@patch
@delegates(subplots)
def show_probas(self:Learner, figsize=(6,6), ds_idx=1, dl=None, one_batch=False, max_n=None, **kwargs):
    recorder = copy(self.recorder) # This is to avoid loss of recorded values while generating preds
    if one_batch: dl = self.dls.one_batch()
    probas, targets = self.get_preds(ds_idx=ds_idx, dl=[dl] if dl is not None else None)
    if probas.ndim == 2 and probas.min() < 0 or probas.max() > 1: probas = nn.Softmax(-1)(probas)
    if not isinstance(targets[0].item(), Integral): return
    targets = targets.flatten()
    if max_n is not None:
        idxs = np.random.choice(len(probas), max_n, False)
        probas, targets = probas[idxs], targets[idxs]
    if isinstance(probas, torch.Tensor): probas = probas.detach().cpu().numpy()
    if isinstance(targets, torch.Tensor): targets = targets.detach().cpu().numpy()
    fig = plt.figure(figsize=figsize, **kwargs)
    classes = np.unique(targets)
    nclasses = len(classes)
    vals = np.linspace(.5, .5 + nclasses - 1, nclasses)[::-1]
    plt.vlines(.5, min(vals) - 1, max(vals), color='black', linewidth=.5)
    cm = plt.get_cmap('gist_rainbow')
    color = [cm(1.* c/nclasses) for c in range(1, nclasses + 1)][::-1]
    # class_probas = np.array([probas[i,t] for i,t in enumerate(targets)])
    class_probas = np.array([probas[i][t] for i,t in enumerate(targets)])
    for i, c in enumerate(classes):
        plt.scatter(class_probas[targets == c] if nclasses > 2 or i > 0 else 1 - class_probas[targets == c],
                    targets[targets == c] + .5 * (np.random.rand((targets == c).sum()) - .5), color=color[i], edgecolor='black', alpha=.2, s=100)
        if nclasses > 2: plt.vlines((targets == c).mean(), i - .5, i + .5, color='r', linewidth=.5)
    plt.hlines(vals, 0, 1)
    plt.ylim(min(vals) - 1, max(vals))
    plt.xlim(0,1)
    plt.xticks(np.linspace(0,1,11), fontsize=12)
    plt.yticks(classes, [self.dls.vocab[x] for x in classes], fontsize=12)
    plt.title('Predicted proba per true class' if nclasses > 2 else 'Predicted class 1 proba per true class', fontsize=14)
    plt.xlabel('Probability', fontsize=12)
    plt.ylabel('True class', fontsize=12)
    plt.grid(axis='x', color='gainsboro', linewidth=.2)
    plt.show()
    self.recorder = recorder

# Cell
@patch
def plot_confusion_matrix(self:Learner, ds_idx=1, dl=None, thr=.5, normalize=False, title='Confusion matrix', cmap="Blues", norm_dec=2, figsize=(6,6),
                          title_fontsize=16, fontsize=12, plot_txt=True, **kwargs):
        "Plot the confusion matrix, with `title` and using `cmap`."
        # This function is mainly copied from the sklearn docs
        assert self.dls.cat
        if self.dls.c == 2: # binary classification
            probas, preds = self.get_preds(ds_idx=ds_idx, dl=dl)
            y_pred = (probas[:, 1] > thr).numpy().astype(int)
            y_test = preds.numpy()
            if normalize: skm_normalize = 'true'
            else: skm_normalize = None
            cm = skm.confusion_matrix(y_test, y_pred, normalize=skm_normalize)
        else:
            cm = ClassificationInterpretation.from_learner(self).confusion_matrix()

        if normalize: cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        fig = plt.figure(figsize=figsize, **kwargs)
        plt.imshow(cm, interpolation='nearest', cmap=cmap)
        if self.dls.c == 2:
            plt.title(f"{title} (threshold: {thr})", fontsize=title_fontsize)
        else:
            plt.title(title, fontsize=title_fontsize)
        tick_marks = np.arange(len(self.dls.vocab))
        plt.xticks(tick_marks, self.dls.vocab, rotation=90, fontsize=fontsize)
        plt.yticks(tick_marks, self.dls.vocab, rotation=0, fontsize=fontsize)

        if plot_txt:
            thresh = cm.max() / 2.
            for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
                coeff = f'{cm[i, j]:.{norm_dec}f}' if normalize else f'{cm[i, j]}'
                plt.text(j, i, coeff, horizontalalignment="center", verticalalignment="center", color="white" if cm[i, j] > thresh else "black", fontsize=fontsize)

        ax = fig.gca()
        ax.set_ylim(len(self.dls.vocab)-.5,-.5)

        plt.tight_layout()
        plt.ylabel('Actual', fontsize=fontsize)
        plt.xlabel('Predicted', fontsize=fontsize)
        plt.grid(False)

# Cell
@patch
def feature_importance(self:Learner, X=None, y=None, partial_n=None, feature_names=None, key_metric_idx=0, show_chart=True, save_df_path=False,
                       random_state=23):
    r"""Calculates feature importance defined to be the change in a model validation loss or metric when a single feature value is randomly shuffled

    This procedure breaks the relationship between the feature and the target, thus the change in the model validation loss or metric is indicative of
    how much the model depends on the feature.

    Args:
        X: array-like object  containing the time series data for which importance will be measured. If None, all data in the validation set will be used.
        y: array-like object containing the targets. If None, all targets in the validation set will be used.
        partial_n: number of samples (if int) or percent of the validation set (if float) that will be used to measure feature importance. If None,
                   all data will be used.
        feature_names (Optional[list(str)]): list of feature names that will be displayed if available. Otherwise they will be var_0, var_1, etc.
        key_metric_idx (Optional[int]): integer to select the metric used in the calculation. If None or no metric is available,
                                        the change is calculated using the validation loss.
        show_chart (bool): flag to indicate if a chart showing permutation feature importance will be plotted.
        save_df_path (str): path to saved dataframe containing the permutation feature importance results.
        random_state (int): controls the shuffling applied to the data. Pass an int for reproducible output across multiple function calls.
    """

    if X is None:
        X = self.dls.valid.dataset.tls[0].items
    if y is None:
        y = self.dls.valid.dataset.tls[1].items
    if partial_n is not None:
        if isinstance(partial_n, float):
            partial_n = int(round(partial_n * len(X)))
        rand_idxs = random_shuffle(np.arange(len(X)), random_state=random_state)[:partial_n]
        X = X.oindex[rand_idxs] if hasattr(X, 'oindex') else X[rand_idxs]
        y = y.oindex[rand_idxs] if hasattr(y, 'oindex') else y[rand_idxs]

    metrics = [mn for mn in self.recorder.metric_names if mn not in ['epoch', 'train_loss', 'valid_loss', 'time']]
    if len(metrics) == 0 or key_metric_idx is None:
        metric_name = self.loss_func.__class__.__name__
        key_metric_idx = None
    else:
        metric_name = metrics[key_metric_idx]
        metric = self.recorder.metrics[key_metric_idx].func
    metric_name = metric_name.replace("train_", "").replace("valid_", "")
    print(f'Selected metric: {metric_name}')

    # Adapted from https://www.kaggle.com/cdeotte/lstm-feature-importance by Chris Deotte (Kaggle GrandMaster)
    if feature_names is None:
        feature_names = [f"var_{i}" for i in range(X.shape[1])]
    else:
        feature_names = listify(feature_names)
    assert len(feature_names) == X.shape[1]

    results = []
    print('Computing feature importance...')

    COLS = ['BASELINE'] + list(feature_names)
    sel_var_ids = np.arange(1, X.shape[1] + 1)[self.dls.sel_vars]

    try:
        for k in progress_bar(range(X.shape[1])):
            if k not in sel_var_ids: continue
            if k>0:
                save_feat = X[:, k-1].copy()
                X[:, k-1] = random_shuffle(X[:, k-1].flatten(), random_state=random_state).reshape(X[:, k-1].shape)
            if key_metric_idx is None:
                value = self.get_X_preds(X, y, with_loss=True)[-1].mean().item()
            else:
                output = self.get_X_preds(X, y)
                value = metric(output[0], output[1]).item()
            print(f"{k:3} feature: {COLS[k]:20} {metric_name}: {value:8.6f}")
            results.append([COLS[k], value])
            del output, value;gc.collect()
            if k>0:
                X[:, k-1] = save_feat
                del save_feat; gc.collect()


    except KeyboardInterrupt:
        if k>0:
            X[:, k-1] = save_feat
            del save_feat; gc.collect()

    # Display feature importance
    if show_chart:
        print()
        df = pd.DataFrame(results, columns=["Feature", metric_name])
        df = df.sort_values(metric_name, ascending=key_metric_idx is None)
        plt.figure(figsize=(10, .5*len(results)))
        plt.barh(np.arange(len(results)), df[metric_name], color="darkblue")
        plt.yticks(np.arange(len(results)), df["Feature"].values)
        plt.title('Permutation Feature Importance', size=16)
        plt.xlabel(f"{metric_name}")
        plt.ylim((-1,len(results)))
        plt.show()

    # Save feature importance
    if save_df_path:
        df = df.sort_values(metric_name,ascending=False)
        df.to_csv(f'{save_df_path}.csv', index=False)