# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0
import torch
import numpy as np
from pytorch_transformers.optimization import (
    # AdamW,
    WarmupConstantSchedule,
    WarmupLinearSchedule,
)

# AdamW(optimizer_grouped_parameters, lr=base_lr, correct_bias=False)

class Optimizer(object):
    def __init__(
        self, 
        parameters, 
        optimizer, 
        lr, 
        lr_scheduler,
        **kwargs):
        
        # Setup torch optimizer
        self.opt_type = optimizer
        self.sch_type = lr_scheduler
        opt = getattr(torch.optim, optimizer)
        if lr_scheduler == 'warmup':
            warmup_step = 4000.0
            self.init_lr = lr
            self.lr_scheduler = []
            if isinstance(self.init_lr, list):
                for lr in self.init_lr:
                    self.lr_scheduler.append(
                        lambda step: lr * warmup_step **0.5 * np.minimum((step+1)*warmup_step**-1.5,(step+1)**-0.5 )
                    )
            else:
                 self.lr_scheduler.append(
                     lambda step: self.init_lr * warmup_step **0.5 * np.minimum((step+1)*warmup_step**-1.5,(step+1)**-0.5 ))
            self.opt = opt(parameters, lr=1.0)
        else:
            self.lr_scheduler = None
            self.opt = opt(parameters, lr=lr)

    def get_opt_state_dict(self):
        return self.opt.state_dict()

    def load_opt_state_dict(self, state_dict):
        self.opt.load_state_dict(state_dict)

    def get_params(self):
        for i, param_group in enumerate(self.opt.param_groups):
            print(param_group)


    def pre_step(self, step):
        if self.lr_scheduler is not None:
            # cur_lr = self.lr_scheduler(step)
            for i, param_group in enumerate(self.opt.param_groups):
                param_group['lr'] = self.lr_scheduler[i](step)
                assert i < len(self.lr_scheduler)
        self.opt.zero_grad()

    def step(self):
        self.opt.step()
