# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/73_callback.captum.ipynb (unless otherwise specified).

__all__ = ['json_clean', 'IntegradedGradientsCallback', 'CaptumInsightsCallback']

# Cell
import tempfile
from ..basics import *
from ..learner import Callback

# Cell

# Dirty hack as json_clean doesn't support CategoryMap type

from ipykernel import jsonutil

_json_clean=jsonutil.json_clean
def json_clean(o):
    o = list(o.items) if isinstance(o,CategoryMap) else o
    return _json_clean(o)

jsonutil.json_clean = json_clean

# Cell
from captum.attr import IntegratedGradients
from captum.attr import visualization as viz

from matplotlib.colors import LinearSegmentedColormap


from captum.insights import AttributionVisualizer, Batch
from captum.insights.features import ImageFeature

# Cell
class IntegradedGradientsCallback(Callback):
    "Captum Callback for Resnet Interpretation"
    def __init__(self):
        pass

    def after_fit(self):
        self.integrated_gradients = IntegratedGradients(self.model)

    def visualize(self, inp_data, n_steps=200, cmap_name='custom blue', colors=None, N=256,
                  methods=None, signs=None, outlier_perc=1):
        if methods is None: methods=['original_image','heat_map']
        if signs is None: signs=["all", "positive"]
        dl = self.dls.test_dl(L(inp_data),with_labels=True, bs=1)
        self.enc_inp,self.enc_preds= dl.one_batch()
        dec_data=dl.decode((self.enc_inp,self.enc_preds))
        self.dec_img,self.dec_pred=dec_data[0][0],dec_data[1][0]
        self.colors = [(0, '#ffffff'),(0.25, '#000000'),(1, '#000000')] if colors is None else colors
        self.attributions_ig = self.integrated_gradients.attribute(self.enc_inp.to(self.dl.device), target=self.enc_preds, n_steps=200)
        default_cmap = LinearSegmentedColormap.from_list(cmap_name,
                                                 self.colors, N=N)
        _ = viz.visualize_image_attr_multiple(np.transpose(self.attributions_ig.squeeze().cpu().detach().numpy(), (1,2,0)),
                             np.transpose(self.dec_img.numpy(), (1,2,0)),
                             methods=methods,
                             cmap=default_cmap,
                             show_colorbar=True,
                             signs=signs,
                             outlier_perc=outlier_perc, titles=[f'Original Image - ({self.dec_pred})', 'IG'])

# Cell
class CaptumInsightsCallback(Callback):
    "Captum Insights Callback for Image Interpretation"
    def __init__(self): pass

    def _formatted_data_iter(self, dl, normalize_func):
        dl_iter=iter(dl)
        while True:
            images,labels=next(dl_iter)
            images=normalize_func.decode(images).to(dl.device)
            yield Batch(inputs=images, labels=labels)

    def visualize(self, inp_data, debug=True):
        _baseline_func= lambda o: o*0
        _get_vocab = lambda vocab: list(map(str,vocab)) if isinstance(vocab[0],bool) else vocab
        dl = self.dls.test_dl(L(inp_data),with_labels=True, bs=4)
        normalize_func= next((func for func in dl.after_batch if type(func)==Normalize),noop)

        visualizer = AttributionVisualizer(
            models=[self.model],
            score_func=lambda o: torch.nn.functional.softmax(o, 1),
            classes=_get_vocab(dl.vocab),
            features=[
                ImageFeature(
                    "Image",
                    baseline_transforms=[_baseline_func],
                    input_transforms=[normalize_func],
                )
            ],
            dataset=self._formatted_data_iter(dl,normalize_func)
        )
        visualizer.render(debug=debug)
