import copy
import torch
import torch.nn as nn


class EMA:
    """Exponential moving average of model weights.

    Maintains a shadow copy of the model. Call update() after each
    optimizer step. Sample from ema.model, never from the training model.
    """

    def __init__(self, model: nn.Module, decay: float = 0.9999):
        self.decay = decay
        self.model = copy.deepcopy(model).eval()
        for p in self.model.parameters():
            p.requires_grad_(False)

    @torch.no_grad()
    def update(self, model: nn.Module) -> None:
        for p_ema, p in zip(self.model.parameters(), model.parameters()):
            p_ema.data.mul_(self.decay).add_(p.data, alpha=1.0 - self.decay)
