# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/utils/12_ema_network.ipynb.

# %% auto 0
__all__ = ['EMANetwork']

# %% ../../nbs/utils/12_ema_network.ipynb 2
from torch import nn
from torch.nn.modules.module import _addindent

# %% ../../nbs/utils/12_ema_network.ipynb 5
class EMANetwork(nn.Module):
    def __init__(self, model: nn.Module, decay: float | object):
        super().__init__()

        self.model = model
        self.decay = decay

    def get_decay(self) -> float:
        if isinstance(self.decay, float):
            return self.decay
        return self.decay.get()

    def update_decay(self):
        if not isinstance(self.decay, float):
            self.decay.step()

    def forward(self, weights: nn.Module | dict) -> nn.Module:
        # Ensure weights is a state_dict
        if isinstance(weights, nn.Module):
            weights = weights.state_dict()

        # Sanity check
        if not set(weights.keys()).issubset(set(self.model.state_dict().keys())):
            raise ValueError("Weights do not match the EMA model's weights")

        # Get decay value
        decay = self.get_decay()

        # Perform EMA
        for name, param in self.model.named_parameters():
            if name not in weights:
                continue
            param.data = decay * param.data + (1 - decay) * weights[name].data

        # Update decay if it's a scheduler
        self.update_decay()

        return self.model

    def __repr__(self):
        return (
            "EMANetwork(\n"
            f"  decay={_addindent(repr(self.decay), 2)}\n"
            f"  model={_addindent(repr(self.model), 2)}\n"
            ")"
        )
