import torch
from torch.optim.adamw import AdamW
import math

class MetaAdamW(AdamW):
    def __init__(self, net, *args, **kwargs):
        super(MetaAdamW, self).__init__(*args, **kwargs)
        self.net = net

    def set_parameter(self, current_module, name, parameters):
        if '.' in name:
            name_split = name.split('.')
            module_name = name_split[0]
            rest_name = '.'.join(name_split[1:])
            for children_name, children in current_module.named_children():
                if module_name == children_name:
                    self.set_parameter(children, rest_name, parameters)
                    break
        else:
            current_module._parameters[name] = parameters

    def meta_step(self, grads):
        group = self.param_groups[0]

        for name, parameter, grad in zip(self.net.get_param_names(), self.net.get_params(), grads):
            if grad is None:
                continue
            
            if grad.is_sparse:
                raise RuntimeError('AdamW does not support sparse gradients')
            
            parameter.detach_()
            # Perform stepweight decay
            parameter = parameter.mul(1 - group['lr'] * group['weight_decay'])
            
            amsgrad = group['amsgrad']
            
            state = self.state[parameter]
            
            # State initialization
            if len(state) == 0:
                state['step'] = 0
                # Exponential moving average of gradient values
                state['exp_avg'] = torch.zeros_like(parameter, memory_format=torch.preserve_format)
                # Exponential moving average of squared gradient values
                state['exp_avg_sq'] = torch.zeros_like(parameter, memory_format=torch.preserve_format)
                if amsgrad:
                    # Maintains max of all exp. moving avg. of sq. grad. values
                    state['max_exp_avg_sq'] = torch.zeros_like(parameter, memory_format=torch.preserve_format)
            
            exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
            if amsgrad:
                max_exp_avg_sq = state['max_exp_avg_sq']
            beta1, beta2 = group['betas']
            
            state['step'] += 1
            bias_correction1 = 1 - beta1 ** state['step']
            bias_correction2 = 1 - beta2 ** state['step']
            
            # Decay the first and second moment running average coefficient
            exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
            exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
            
            # Deal with stability issues
            mask = exp_avg_sq == 0.
            _maybe_mask(exp_avg_sq, mask)
            
            if amsgrad:
                # Maintains the maximum of all 2nd moment running avg. till now
                torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
                # Use the max. for normalizing running avg. of gradient
                denom = (max_exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
            else:
                denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])

            step_size = group['lr'] / bias_correction1
            
            self.set_parameter(self.net, name, parameter.addcdiv(exp_avg, denom, value=-step_size))
            
def _get_mask_closure(mask):
    def closure(grad):
        grad = torch.where(mask, torch.zeros_like(grad), grad)
        if grad.requires_grad:
            grad.register_hook(_get_mask_closure(mask))
        return grad
    return closure


def _maybe_mask(tensor, mask):
    if tensor.requires_grad:
        tensor.register_hook(_get_mask_closure(mask))