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

__all__ = ['ClassificationInterpretationEx']

# Cell
try:
    from fastai.vision.all import *
    from fastai.metrics import *
except:
    from fastai2.vision.all import *
    from fastai2.metrics import *

# Cell
class ClassificationInterpretationEx(ClassificationInterpretation):
    """
    Extend fastai2's `ClassificationInterpretation` to analyse model predictions in more depth
    """
    def plot_confusion_matrix(interp, normalize=False, title='Confusion matrix', cmap="Blues", norm_dec=2,
                              plot_txt=True, return_fig=False, **kwargs):
        """
        Plot the confusion matrix, with `title` and using `cmap`.
        An exact replica of fastai2's method, with the added option
        of `return_fig`, to be able to save the image to disk
        """
        # This function is mainly copied from the sklearn docs
        cm = interp.confusion_matrix()
        if normalize: cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        fig = plt.figure(**kwargs)
        plt.imshow(cm, interpolation='nearest', cmap=cmap)
        plt.title(title)
        tick_marks = np.arange(len(interp.vocab))
        plt.xticks(tick_marks, interp.vocab, rotation=90)
        plt.yticks(tick_marks, interp.vocab, rotation=0)

        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")

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

        plt.tight_layout()
        plt.ylabel('Actual')
        plt.xlabel('Predicted')
        plt.grid(False)
        if return_fig: return fig

    def compute_label_confidence(self):
        """
        Collate prediction confidence, filenames, and ground truth labels
        in DataFrames, and store them as class attributes
        `self.preds_df` and `self.preds_df_each`
        """
        self._preds_collated = [
            (item, self.dl.vocab[label_idx], *preds.numpy()*100)\
            for item,label_idx,preds in zip(self.dl.items,
                                            self.targs,
                                            self.preds)
        ]

        self.preds_df       = pd.DataFrame(self._preds_collated, columns = ['fname','truth', *self.dl.vocab])
        self._preds_df_each = {l:self.preds_df.copy()[self.preds_df.truth == l].reset_index(drop=True) for l in self.dl.vocab}
        self.preds_df_each  = defaultdict(dict)


        sort_desc = lambda x,col: x.sort_values(col, ascending=False).reset_index(drop=True)
        for label,df in self._preds_df_each.items():
            filt = df[label] == df[self.dl.vocab].max(axis=1)
            self.preds_df_each[label]['accurate']   = df.copy()[filt]
            self.preds_df_each[label]['inaccurate'] = df.copy()[~filt]

            self.preds_df_each[label]['accurate']   = sort_desc(self.preds_df_each[label]['accurate'], label)
            self.preds_df_each[label]['inaccurate'] = sort_desc(self.preds_df_each[label]['inaccurate'], label)
            assert len(self.preds_df_each[label]['accurate']) + len(self.preds_df_each[label]['inaccurate']) == len(df)

    def plot_label_confidence(self, bins:int=10, fig_width:int=12, fig_height_base:int=4,
                              title:str='Accurate vs. Inaccurate Predictions Confidence (%) Levels Per Label',
                              return_fig:bool=False, label_bins:bool=True,
                              accurate_color='mediumseagreen', inaccurate_color='tomato'):
        'Plot label confidence histograms for each label'
        if not hasattr(self, 'preds_df_each'): self.compute_label_confidence()
        fig, axes = plt.subplots(nrows = len(self.preds_df_each.keys()), ncols=2,
                                 figsize = (fig_width, fig_height_base * len(self.dl.vocab)))
        for i, (label, df) in enumerate(self.preds_df_each.items()):
            height=0
            # find max height
            for mode in ['inaccurate', 'accurate']:
                len_bins,_ = np.histogram(df[mode][label], bins=10)
                if len_bins.max() > height: height=len_bins.max()

            for mode,ax in zip(['inaccurate', 'accurate'], axes[i]):
                range_ = (50,100) if mode == 'accurate' else (0,50)
                color  = accurate_color if mode == 'accurate' else inaccurate_color
                num,_,patches = ax.hist(df[mode][label], bins=bins, range=range_, rwidth=.95, color=color)
                if label_bins:
                    for rect in patches:
                        ht = rect.get_height()
                        ax.annotate(s  = f"{int(ht) if ht > 0 else ''}",
                            xy = (rect.get_x() + rect.get_width()/2, ht),
                            xytext = (0,3), # offset vertically by 3 points
                            textcoords = 'offset points',
                            ha = 'center', va = 'bottom'
                           )
                ax.set_ybound(upper=height + height*0.3)
                ax.set_xlabel(f'{label}: {mode.capitalize()}')
                ax.set_ylabel(f'No. {mode.capitalize()} = {len(df[mode][label])}')
        fig.suptitle(title)
        plt.subplots_adjust(top = 0.9, bottom=0.01, hspace=0.25, wspace=0.2)
        if return_fig: return fig


    def plot_accuracy(self, width=0.5, figsize=(6,6), return_fig=False,
                      title='Accuracy Per Label', ylabel='Accuracy (%)',
                      color='steelblue', vertical_labels=False):
        'Plot a bar plot showing accuracy per label'
        if not hasattr(self, 'preds_df_each'): self.compute_label_confidence()
        self.accuracy_dict = defaultdict()

        for label,df in self.preds_df_each.items():
            total = len(df['accurate']) + len(df['inaccurate'])
            self.accuracy_dict[label] = 100 * len(df['accurate']) / total

        fig,ax = plt.subplots(figsize=figsize)

        x = self.accuracy_dict.keys()
        y = [v for k,v in self.accuracy_dict.items()]

        rects = ax.bar(x,y,width,color=color)
        for rect in rects:
            ht = rect.get_height()
            ax.annotate(s  = f"{ht:.02f}",
                        xy = (rect.get_x() + rect.get_width()/2, ht),
                        xytext = (0,3), # offset vertically by 3 points
                        textcoords = 'offset points',
                        ha = 'center', va = 'bottom'
                       )
        ax.set_ybound(lower=0, upper=100)
        ax.set_yticks(np.arange(0,110,10))
        ax.set_ylabel(ylabel)
        ax.set_xticklabels(x, rotation='vertical' if vertical_labels else 'horizontal')
        plt.suptitle(title)
        if return_fig: return fig

    def get_fnames(self, label:str,
                   mode:('accurate','inaccurate'),
                   conf_level:Union[int,float,tuple]) -> np.ndarray:
        """
        Utility function to grab filenames of a particular label `label` that were classified
        as per `mode` (accurate|inaccurate).
        These filenames are filtered by `conf_level` which can be above or below a certain
        threshold (above if `mode` == 'accurate' else below), or in confidence ranges
        """
        assert label in self.dl.vocab
        if not hasattr(self, 'preds_df_each'): self.compute_label_confidence()
        df = self.preds_df_each[label][mode].copy()
        if mode == 'accurate':
            if isinstance(conf_level, tuple):       filt = df[label].between(*conf_level)
            if isinstance(conf_level, (int,float)): filt = df[label] > conf_level
        if mode == 'inaccurate':
            if isinstance(conf_level, tuple):       filt = df[label].between(*conf_level)
            if isinstance(conf_level, (int,float)): filt = df[label] < conf_level
        return df[filt].fname.values