<reponame>brudermueller/OpenPCDet
# This file is modified from https://github.com/traveller59/second.pytorch

import math
import os
import numpy as np 
from functools import partial

import numpy as np
import torch.optim.lr_scheduler as lr_sched

# from .fastai_optim import OptimWrapper
from fastai_optim import OptimWrapper
# from . import OptimWrapper
# 
# from tools.train_utils.fastai_otim import OptimWrapper


class LRSchedulerStep(object):
    def __init__(self, fai_optimizer: OptimWrapper, total_step, lr_phases,
                 mom_phases):
        # if not isinstance(fai_optimizer, OptimWrapper):
        #     raise TypeError('{} is not a fastai OptimWrapper'.format(
        #         type(fai_optimizer).__name__))
        self.optimizer = fai_optimizer
        self.total_step = total_step
        self.lr_phases = []

        for i, (start, lambda_func) in enumerate(lr_phases):
            if len(self.lr_phases) != 0:
                assert self.lr_phases[-1][0] < start
            if isinstance(lambda_func, str):
                lambda_func = eval(lambda_func)
            if i < len(lr_phases) - 1:
                self.lr_phases.append((int(start * total_step), int(lr_phases[i + 1][0] * total_step), lambda_func))
            else:
                self.lr_phases.append((int(start * total_step), total_step, lambda_func))
        assert self.lr_phases[0][0] == 0
        self.mom_phases = []
        for i, (start, lambda_func) in enumerate(mom_phases):
            if len(self.mom_phases) != 0:
                assert self.mom_phases[-1][0] < start
            if isinstance(lambda_func, str):
                lambda_func = eval(lambda_func)
            if i < len(mom_phases) - 1:
                self.mom_phases.append((int(start * total_step), int(mom_phases[i + 1][0] * total_step), lambda_func))
            else:
                self.mom_phases.append((int(start * total_step), total_step, lambda_func))
        assert self.mom_phases[0][0] == 0

    def step(self, step):
        for start, end, func in self.lr_phases:
            if step >= start:
                self.optimizer.lr = func((step - start) / (end - start))
        for start, end, func in self.mom_phases:
            if step >= start:
                self.optimizer.mom = func((step - start) / (end - start))


def annealing_cos(start, end, pct):
    # print(pct, start, end)
    "Cosine anneal from `start` to `end` as pct goes from 0.0 to 1.0."
    cos_out = np.cos(np.pi * pct) + 1
    return end + (start - end) / 2 * cos_out


class OneCycle(LRSchedulerStep):
    def __init__(self, fai_optimizer, total_step, lr_max, moms, div_factor,
                 pct_start):
        self.lr_max = lr_max
        self.moms = moms
        self.div_factor = div_factor
        self.pct_start = pct_start
        a1 = int(total_step * self.pct_start)
        a2 = total_step - a1
        low_lr = self.lr_max / self.div_factor
        lr_phases = ((0, partial(annealing_cos, low_lr, self.lr_max)),
                     (self.pct_start,
                      partial(annealing_cos, self.lr_max, low_lr / 1e4)))
        mom_phases = ((0, partial(annealing_cos, *self.moms)),
                      (self.pct_start, partial(annealing_cos,
                                               *self.moms[::-1])))
        fai_optimizer.lr, fai_optimizer.mom = low_lr, self.moms[0]
        super().__init__(fai_optimizer, total_step, lr_phases, mom_phases)


class CosineWarmupLR(lr_sched._LRScheduler):
    def __init__(self, optimizer, T_max, eta_min=0, last_epoch=-1):
        self.T_max = T_max
        self.eta_min = eta_min
        super(CosineWarmupLR, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        return [self.eta_min + (base_lr - self.eta_min) *
                (1 - math.cos(math.pi * self.last_epoch / self.T_max)) / 2
                for base_lr in self.base_lrs]


class FakeOptim:
    def __init__(self):
        self.lr = 0
        self.mom = 0


if __name__ == "__main__":
#     import matplotlib.pyplot as plt

#     opt = FakeOptim()  # 3e-3, wd=0.4, div_factor=10
#     # schd = OneCycle(opt, 100, 3e-3, (0.95, 0.85), 10.0, 0.1)
#     # total_step_size = 108720
#     # epochs = [30, 50, 80]
#     num_epochs = 30 
#     total_step_size = num_epochs * 3624
#     schd = OneCycle(opt, total_step_size, 0.01, (0.95, 0.85), 100.0, 0.4)


#     lrs = []
#     moms = []
#     for i in range(total_step_size):
#         schd.step(i)
#         lrs.append(opt.lr)
#         moms.append(opt.mom)
#     plot_dir = os.path.dirname(os.path.abspath(__file__))
#     plot_dir = os.path.join(plot_dir, "../../../plots/") 
    

#     plt.xlabel('Iterations', fontsize=16)
#     plt.ylabel('Learning rate',fontsize=16)
#     plt.xticks(fontsize=14)
#     plt.yticks(fontsize=14)
#     plt.plot(lrs, label="30")
#     plt.legend(loc='upper right', title='# epochs', frameon=True, fontsize='medium')
#     plt.vlines(total_step_size/2, 0.0, 0.01, color='gray', alpha=0.3, linestyles='dashed')
#     plt.savefig(plot_dir+'lrs_1.png', bbox_inches='tight')  
#     # plt.show()
#     plt.plot(moms)
#     plt.savefig(plot_dir+'moms_1.png', bbox_inches='tight')  
#     # plt.show()
    COLORS = [
        '#1f77b4', '#aec7e8', '#ff7f0e', '#ffbb78', '#2ca02c',
        '#98df8a', '#d62728', '#ff9896', '#9467bd', '#c5b0d5',
        '#8c564b', '#c49c94', '#e377c2', '#f7b6d2', '#7f7f7f',
        '#c7c7c7', '#bcbd22', '#dbdb8d', '#17becf', '#9edae5']
    import matplotlib.pyplot as plt

    opt = FakeOptim()  # 3e-3, wd=0.4, div_factor=10
    # schd = OneCycle(opt, 100, 3e-3, (0.95, 0.85), 10.0, 0.1)
    # total_step_size = 108720
    # epochs = [30, 50, 80] 
    plot_dir = os.path.dirname(os.path.abspath(__file__))
    plot_dir = os.path.join(plot_dir, "../../../plots/") 
    num_epochs = 30 
    total_step_size = num_epochs * 3624

    wd = [0.3, 0.4]#[0.3, 0.4, 0.5]
    div_factor = [10, 100]
    max_lr = 0.01

    color_ct = 0 
    for pct_start in wd: 
        for div in div_factor: 
            start_lr = max_lr/div
            color_ct += 1
            schd = OneCycle(opt, total_step_size, max_lr, (0.95, 0.85), div, pct_start)
            lrs = []
            moms = []
            for i in range(total_step_size):
                schd.step(i)
                lrs.append(opt.lr)
                moms.append(opt.mom)

            plt.xlabel('Iterations', fontsize=16)
            plt.ylabel('Learning rate',fontsize=16)
            plt.xticks(fontsize=14)
            plt.yticks(fontsize=14)
            plt.plot(lrs, label="rise: {}\nstart_lr: {}".format(pct_start, start_lr), color=COLORS[color_ct])
            plt.legend(loc='upper right', title='Parameters', frameon=True, fontsize='medium')
            plt.vlines(total_step_size/2, 0.0, 0.01, color='gray', alpha=0.3, linestyles='dashed')
    plt.savefig(plot_dir+'lrs_2.pdf', bbox_inches='tight')  
    # plt.show()
    plt.plot(moms)
    plt.savefig(plot_dir+'moms_1.png', bbox_inches='tight')  