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

__all__ = ['create_classifier', 'Learner', 'InterpretModel']

# Cell
import tensorflow as tf
import tensorflow_addons as tfa
from tensorflow.keras.models import Model
from .datagenerator import Dataset

from typeguard import check_argument_types, check_return_type

# Cell
from PIL import Image
import numpy as np
from tf_keras_vis.utils import normalize
from tf_keras_vis.gradcam import GradcamPlusPlus, Gradcam
from functools import partial

import matplotlib.pyplot as plt
import matplotlib.cm as cm

# Cell
def create_classifier(base_model_fn:callable, num_classes:int,
                      weights='imagenet', dropout=0,
                      include_top=False,
                      name=None):

    outputs = 1 if num_classes == 2 else num_classes

    base_model = base_model_fn(
        include_top=include_top,
        weights=weights,
    )
    if include_top: return base_model
    model = tf.keras.Sequential(name=name)
    model.add(base_model)
    model.add(tf.keras.layers.GlobalAveragePooling2D())
    model.add(tf.keras.layers.Dropout(dropout))
    model.add(tf.keras.layers.Dense(outputs, name='output'))

    return model

# Cell
class Learner(Model):
    """
    The Learner class inherits tf.keras.Model and contains everything a model needs for training.
    It exposes learner.cyclic_fit method which trains the model using Cyclic Learning rate discovered by Leslie Smith.

    Arguments:
    ds: Dataset object
    base_model: callable function from keras.applications

    """
    _AUTOTUNE = tf.data.experimental.AUTOTUNE

    def __init__(self, ds: Dataset, base_model_fn:callable, pretrained:bool=True, include_top=False, **kwargs):
        assert check_argument_types()

        super(Learner, self).__init__()
        self.ds = ds
        self.total = len(ds)
        self.NUM_CLASSES = ds.NUM_CLASSES
        self.gradcam = None
        self.include_top = include_top

        weights = 'imagenet' if pretrained else None

        #self.base_model = base_model_fn.name

        self.model = create_classifier(
            base_model_fn,
            self.NUM_CLASSES,
            dropout=kwargs.get('dropout', 0.5),
            include_top=include_top,
            name=kwargs.get('name', None)
        )


    def build(self): pass

    def summary(self): return self.model.summary()

    #def get_layer(name=None, index=None): return self.model(name, index)

    def compile(self, *args, **kwargs): return self.model.compile(*args, **kwargs)

    def call(self, *args, **kwargs): return self.model.call(*args, **kwargs)

    def fit(self, *args, **kwargs): return self.model.fit(*args, **kwargs)

    def warmup(self):pass

    def prewhiten(self, image):
        image = tf.cast(image, tf.float32)
        image = image / 127.5 - 1.0
        return image

    def rescale(self, image, label):
        image = self.prewhiten(image)
        return image, label

    def _get_optimizer(self,
                       optimizer,
                       momentum=0.9,
                       **kwargs
                      ):
        if optimizer.__name__=='SGD':
            optimizer = partial(optimizer,
                momentum=momentum,
                nesterov=kwargs.get('nesterov', True)
            )
        else:
            optimizer = partial(optimizer,
                momentum=momentum,
            )
        return optimizer


    def _prepare_dl(self, bs=8, shuffle=True):
        ds = self.ds
        dl = ds.get_tf_dataset(shuffle=shuffle)
        dl = dl.map(self.rescale, Learner._AUTOTUNE)
        return dl.batch(bs).prefetch(Learner._AUTOTUNE)


    def cyclic_fit(self,
                   epochs,
                   batch_size,
                   lr_range=(1e-4, 1e-2),
                   optimizer=tf.keras.optimizers.SGD,
                   momentum=0.9,
                   validation_data=None,
                   callbacks=None,
                   *args,
                   **kwargs
                  ):
        """Trains model on ds as train data with cyclic learning rate.
        Dataset will be automatically converted into `tf.data` format and images will be prewhitened in range of [-1, 1].
        """

        self.max_lr, self.min_lr = lr_range
        ds = self.ds

        step_size = 2 * len(self.ds)//batch_size

        lr_schedule = tfa.optimizers.Triangular2CyclicalLearningRate(
                                    initial_learning_rate=lr_range[0],
                                    maximal_learning_rate=lr_range[1],
                                    step_size=kwargs.get('step_size', step_size),
                                    scale_mode=kwargs.get('scale', 'cycle'))


        optimizer = self._get_optimizer(optimizer, momentum=momentum)
        optimizer = optimizer(learning_rate=lr_schedule)
        self.model.optimizer = optimizer

        return self.model.fit(
            self._prepare_dl(batch_size, kwargs.get('shuffle', True)),
            validation_data=validation_data,
            epochs=epochs,
            callbacks=callbacks
        )

    def gradcam(self, image:Image.Image, image_size=None):
        # assert check_argument_types()

        def model_modifier(m):
            """Converts sigmoid to linear
            """
            m.layers[-1].activation = tf.keras.activations.linear
            return m


        if image_size:
            image_size = self.ds.img_sz_list.get_size()
        image = image.resize(image_size)

        X = np.asarray(image, np.float32)
        X = self.prewhiten(X)
        X = np.expand_dims(X, 0)

        if self.gradcam is None:
            self.gradcam = Gradcam(self.model,
                              model_modifier,
                              clone=True)

        cam = gradcam(get_loss,
              X,
              penultimate_layer=-1, # model.layers number
              seek_penultimate_conv_layer=False,

             )


# Cell
class InterpretModel():
    def __init__(self, gradcam_pp:bool, learner:Learner, clone:bool=False):
        """Args:
        gradcam_pp: if True GradCam class will be used else GradCamPlusplus
        clone: whether GradCam will clone learner.model
        """
        if gradcam_pp:
            self.gradcam_fn = GradcamPlusPlus
        else:
            self.gradcam_fn = Gradcam
        self.learner = learner

        self.gradcam = self.gradcam_fn(learner.model,
                                      self.model_modifier,
                                      clone=clone)

        if self.learner.include_top is not True:
            self.gradcam._find_penultimate_output = self.patch


    def __call__(self, image:Image.Image, auto_resize:bool=True, image_size=None):
        #assert check_argument_types()
        gradcam = self.gradcam
        get_loss = self.get_loss
        if auto_resize and image_size is None:
            image_size = self.learner.ds.img_sz_list.get_size()
        if image_size:
            image = image.resize(image_size)

        X = np.asarray(image, np.float32)
        X = self.learner.prewhiten(X)
        X = np.expand_dims(X, 0)

        cam = gradcam(get_loss,
                      X,
                      penultimate_layer=-1, # model.layers number
                      seek_penultimate_conv_layer=True,
                     )
        cam = normalize(cam)
        heatmap = np.uint8(cm.jet(cam[0])[..., :3] * 255)
        plt.imshow(image)
        plt.imshow(heatmap, cmap='jet', alpha=0.5)
        plt.show()


    def patch(self, *args, **kwargs):
        """Path _find_penultimate_output method of tf_keras_vis"""
        if self.learner.include_top:
            return self.learner.model.layers[-1].output
        return self.learner.model.layers[0].get_output_at(-1)


    def model_modifier(self, m):
        """Sets last activation to linear
        """
        m.layers[-1].activation = tf.keras.activations.linear
        return m

    def get_loss(self, preds):
        if self.learner.NUM_CLASSES == 2:
            ret = preds[0]
        else:
            index = tf.argmax(tf.math.softmax(preds), axis=1)[0]
            # print(index, preds.shape)
            ret = preds[0, index]
            print(f'index: {index}')
        return ret