# From https://www.kaggle.com/robotdreams/one-cycle-policy-with-keras

from sklearn.metrics import log_loss, roc_auc_score, accuracy_score
from keras.losses import binary_crossentropy
from keras.metrics import binary_accuracy
from keras import backend as K
from keras.callbacks import *

class OneCycleLR(Callback):

    def __init__(self, base_lr=0.001, max_lr=0.006, step_size=2000.,
                 base_m=.85, max_m=.95, cyclical_momentum=False):

        self.base_lr = base_lr
        self.max_lr = max_lr
        self.base_m = base_m
        self.max_m = max_m
        self.cyclical_momentum = cyclical_momentum
        self.step_size = step_size

        self.clr_iterations = 0.
        self.cm_iterations = 0.
        self.trn_iterations = 0.
        self.history = {}

    def clr(self):

        cycle = np.floor(1+self.clr_iterations/(2*self.step_size))

        if cycle == 2:
            x = np.abs(self.clr_iterations/self.step_size - 2*cycle + 1)
            return self.base_lr-(self.base_lr-self.base_lr/100)*np.maximum(0,(1-x))

        else:
            x = np.abs(self.clr_iterations/self.step_size - 2*cycle + 1)
            return self.base_lr + (self.max_lr-self.base_lr)*np.maximum(0,(1-x))

    def cm(self):

        cycle = np.floor(1+self.clr_iterations/(2*self.step_size))

        if cycle == 2:

            x = np.abs(self.clr_iterations/self.step_size - 2*cycle + 1)
            return self.max_m

        else:
            x = np.abs(self.clr_iterations/self.step_size - 2*cycle + 1)
            return self.max_m - (self.max_m-self.base_m)*np.maximum(0,(1-x))


    def on_train_begin(self, logs={}):
        logs = logs or {}

        if self.clr_iterations == 0:
            K.set_value(self.model.optimizer.lr, self.base_lr)
        else:
            K.set_value(self.model.optimizer.lr, self.clr())

        if self.cyclical_momentum == True:
            if self.clr_iterations == 0:
                K.set_value(self.model.optimizer.momentum, self.cm())
            else:
                K.set_value(self.model.optimizer.momentum, self.cm())


    def on_batch_begin(self, batch, logs=None):

        logs = logs or {}
        self.trn_iterations += 1
        self.clr_iterations += 1

        self.history.setdefault('lr', []).append(K.get_value(self.model.optimizer.lr))
        self.history.setdefault('iterations', []).append(self.trn_iterations)

        if self.cyclical_momentum == True:
            self.history.setdefault('momentum', []).append(K.get_value(self.model.optimizer.momentum))

        for k, v in logs.items():
            self.history.setdefault(k, []).append(v)

        K.set_value(self.model.optimizer.lr, self.clr())

        if self.cyclical_momentum == True:
            K.set_value(self.model.optimizer.momentum, self.cm())
