import itertools
from torch.nn.utils import clip_grad_norm_
from torch.optim import Adam, AdamW


class AdamWithClip(Adam):

    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=0, amsgrad=False, max_norm=None, norm_type=2):
        super(AdamWithClip, self).__init__(params, lr, betas, eps, weight_decay, amsgrad)
        self.max_norm = max_norm
        self.norm_type = norm_type

    def step(self, closure=None):
        if self.max_norm is not None:
            for group in self.param_groups:
                clip_grad_norm_(group['params'], self.max_norm, self.norm_type)
        super(AdamWithClip, self).step(closure)


class AdamWWithClip(AdamW):

    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=0, amsgrad=False, max_norm=None, norm_type=2):
        super(AdamWWithClip, self).__init__(params, lr, betas, eps, weight_decay, amsgrad)
        self.max_norm = max_norm
        self.norm_type = norm_type

    def step(self, closure=None):

        if self.max_norm is not None:
            for group in self.param_groups:
                clip_grad_norm_(group['params'], self.max_norm, self.norm_type)
        super(AdamWWithClip, self).step(closure)


# class AdamWWithClipDev(AdamW):
#     def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
#                  weight_decay=1e-2, amsgrad=False, clip_norm=None, norm_type=2):
#         super(AdamWWithClipDev, self).__init__(params, lr, betas, eps, weight_decay, amsgrad)
#         self.clip_norm = clip_norm
#         self.norm_type = norm_type
#
#     def step(self, closure=None):
#         if self.clip_norm is not None:
#             all_params = itertools.chain(*[x["params"] for x in self.param_groups])
#             clip_grad_norm_(all_params, self.clip_norm, self.norm_type)
#
#         super(AdamWWithClipDev, self).step(closure)

class AdamWWithBackboneClipDev(AdamW):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=1e-2, amsgrad=False, clip_norm=None, norm_type=2):
        super(AdamWWithBackboneClipDev, self).__init__(params, lr, betas, eps, weight_decay, amsgrad)
        self.clip_norm = clip_norm
        self.norm_type = norm_type

    def step(self, closure=None):
        if self.clip_norm is not None:
            all_params = itertools.chain(*[x["params"] for x in self.param_groups if x['params'][0].backbone_specific ])
            clip_grad_norm_(all_params, self.clip_norm, self.norm_type)

        super(AdamWWithBackboneClipDev, self).step(closure)

class AdamWWithClipDev(AdamW):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=1e-2, amsgrad=False, clip_norm=None, norm_type=2):
        super(AdamWWithClipDev, self).__init__(params, lr, betas, eps, weight_decay, amsgrad)
        self.clip_norm = clip_norm
        self.norm_type = norm_type

        self._split_param_groups = None
        self.reset_split_param_groups()

    def reset_split_param_groups(self):
        if self.clip_norm is not None:
            backbone_param, neck_param, decoder_param, task_param = [], [], [], []
            for x in self.param_groups:
                if x["params"][0].backbone_specific:
                    backbone_param.append(x["params"])
                elif x["params"][0].neck_specific:
                    neck_param.append(x["params"])
                elif x["params"][0].decoder_specific:
                    decoder_param.append(x["params"])
                elif x["params"][0].task_specific:
                    task_param.append(x["params"])
            self._split_param_groups = [_g for _g in [backbone_param,
                                                      neck_param,
                                                      decoder_param,
                                                      task_param] if len(_g) > 0]
            print(f">>> reset_split_param_groups, backbone_param: {len(backbone_param)}"
                  f", neck_param: {len(neck_param)}, decoder_param: {len(decoder_param)}"
                  f", task_param: {len(task_param)}")

    def step(self, closure=None):
        if self.clip_norm is not None:
            for _g in self._split_param_groups:
                all_params = itertools.chain(*_g)
                clip_grad_norm_(all_params, self.clip_norm, self.norm_type)

        super(AdamWWithClipDev, self).step(closure)