import torch                                                                                                         


class ExponentialMovingAverageV1(torch.optim.swa_utils.AveragedModel):
    """Maintains moving averages of model parameters using an exponential decay.
    ``ema_avg = decay * avg_model_param + (1 - decay) * model_param``
    `torch.optim.swa_utils.AveragedModel <https://pytorch.org/docs/stable/optim.html#custom-averaging-strategies>`_    is used to compute the EMA.
    """
    def __init__(self, model, decay, device='cpu'):
        ema_avg = (lambda avg_model_param, model_param, num_averaged:
                   decay * avg_model_param + (1 - decay) * model_param)
        super().__init__(model, device, ema_avg) 

    def update_parameters(self, model):
        for p_swa, p_model in zip(self.module.state_dict().values(), model.state_dict().values()):
            device = p_swa.device 
            p_model_ = p_model.detach().to(device)
            if self.n_averaged == 0:
                p_swa.detach().copy_(p_model_)
            else:
                p_swa.detach().copy_(self.avg_fn(p_swa.detach(), p_model_,
                                     self.n_averaged.to(device)))
        self.n_averaged += 1

class ExponentialMovingAverage(torch.optim.swa_utils.AveragedModel):
    """Maintains moving averages of model parameters using an exponential decay.
    ``ema_avg = decay * avg_model_param + (1 - decay) * model_param``
    `torch.optim.swa_utils.AveragedModel <https://pytorch.org/docs/stable/optim.html#custom-averaging-strategies>`_
    is used to compute the EMA.    """
    def __init__(self, model, decay, device='cpu'):
        ema_avg = (lambda avg_model_param, model_param, num_averaged:
                   decay * avg_model_param + (1 - decay) * model_param)
        super().__init__(model, device, ema_avg, use_buffers=True)



class ToyModel(torch.nn.Module):
    
    def __init__(self):
        super().__init__()
        self.x = torch.nn.Parameter(torch.zeros(5))
        self.register_buffer('y', torch.zeros(5))


    def forward(self, input):
        self.x += input
        self.y += input
        return self.x, self.y

decay = 0.9
model1 = ToyModel()
ema1 = ExponentialMovingAverageV1(model1, decay)
model2 = ToyModel()
ema2 = ExponentialMovingAverage(model2, decay)

x = torch.ones(5)

for _ in range(10):
    with torch.no_grad():
        model1(x)
        model2(x)
        ema1.update_parameters(model1)
        ema2.update_parameters(model2)

        assert torch.equal(ema1.module.x, ema2.module.x)
        assert torch.equal(ema1.module.y, ema2.module.y)