import logging

import torch
from torch import Tensor, nn

logger = logging.getLogger(__name__)


class Normalizer(nn.Module):
    def __init__(self, momentum=0.01, eps=1e-9):
        super().__init__()
        self.momentum = momentum
        self.eps = eps
        self.running_mean_unsafe: Tensor
        self.running_var_unsafe: Tensor
        self.register_buffer("running_mean_unsafe", torch.full([], torch.nan))
        self.register_buffer("running_var_unsafe", torch.full([], torch.nan))

    @property
    def started(self):
        return not torch.isnan(self.running_mean_unsafe)

    @property
    def running_mean(self):
        if not self.started:
            return torch.zeros_like(self.running_mean_unsafe)
        return self.running_mean_unsafe

    @property
    def running_std(self):
        if not self.started:
            return torch.ones_like(self.running_var_unsafe)
        return (self.running_var_unsafe + self.eps).sqrt()

    @torch.no_grad()
    def _ema(self, a: Tensor, x: Tensor):
        return (1 - self.momentum) * a + self.momentum * x

    def update_(self, x):
        if not self.started:
            self.running_mean_unsafe = x.mean()
            self.running_var_unsafe = x.var()
        else:
            self.running_mean_unsafe = self._ema(self.running_mean_unsafe, x.mean())
            self.running_var_unsafe = self._ema(self.running_var_unsafe, (x - self.running_mean).pow(2).mean())

    def forward(self, x: Tensor, update=True):
        if self.training and update:
            self.update_(x)
        self.stats = dict(mean=self.running_mean.item(), std=self.running_std.item())
        x = (x - self.running_mean) / self.running_std
        return x

    def inverse(self, x: Tensor):
        return x * self.running_std + self.running_mean