# -*- coding: utf-8 -*-

import os

from keras.callbacks import Callback

class SWA(Callback):

    """

    This callback implements a stochastic weight averaging (SWA) method with constant lr

    as presented in the paper -

        "Izmailov et al. Averaging Weights Leads to Wider Optima and Better Generalization"

        (https://arxiv.org/abs/1803.05407)

    Author's implementation: https://github.com/timgaripov/swa

    """

    def __init__(self, swa_model, checkpoint_dir, model_name, swa_start=1):

        """

        :param swa_model: the model that we use to store the average of the weights once SWA begins

        :param checkpoint_dir: the directory where the model will be saved in

        :param model_name: the name of model we're training

        :param swa_start: the epoch when averaging begins. We generally pre-train the network for

                          a certain amount of epochs to start (swa_start > 1), as opposed to

                          starting to track the average from the very beginning.

        """

        super(SWA, self).__init__()

        self.checkpoint_dir = checkpoint_dir

        self.model_name = model_name

        self.swa_start = swa_start

        self.swa_model = swa_model

    def on_train_begin(self, logs=None):

        self.epoch = 0

        self.swa_n = 0

        '''Note: I found deep copy of a model with customized layer would give errors'''

        # self.swa_model = copy.deepcopy(self.model)  # make a copy of the model we're training

        '''Note: Something wired still happen even though i use keras.models.clone_model method, 

                 so I build swa_model outside this callback and pass it as an argument. 

                 It's not fancy, but the best I can do :)

        '''

        # self.swa_model = keras.models.clone_model(self.model)

        # see: https://github.com/keras-team/keras/issues/1765

        self.swa_model.set_weights(self.model.get_weights())

    def on_epoch_end(self, epoch, logs=None):

        if (self.epoch + 1) >= self.swa_start:

            self.update_average_model()

            self.swa_n += 1

        self.epoch += 1

    def update_average_model(self):

        # update running average of parameters

        alpha = 1. / (self.swa_n + 1)

        for layer, swa_layer in zip(self.model.layers, self.swa_model.layers):

            weights = []

            for w1, w2 in zip(swa_layer.get_weights(), layer.get_weights()):

                weights.append((1 - alpha) * w1 + alpha * w2)

            swa_layer.set_weights(weights)

    def on_train_end(self, logs=None):

        print(f'Logging Info - Saving SWA model checkpoint: {self.model_name}_swa.hdf5')

        self.swa_model.save_weights(os.path.join(self.checkpoint_dir,

                                                 f'{self.model_name}_swa.hdf5'))

        print('Logging Info - SWA model Saved')