"""Shared normalization utilities for ML agents."""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from numpy.typing import NDArray


class RunningMeanStd:
    """Running mean and standard deviation for observation normalization.

    Uses Welford's online algorithm for numerically stable computation
    of running statistics.

    Args:
        shape: Shape of the observations to normalize
        epsilon: Small constant to avoid division by zero

    Example:
        normalizer = RunningMeanStd((328,))
        normalizer.update(observations)
        normalized = normalizer.normalize(obs)
    """

    def __init__(self, shape: tuple[int, ...], epsilon: float = 1e-8) -> None:
        self.mean = np.zeros(shape, dtype=np.float64)
        self.var = np.ones(shape, dtype=np.float64)
        self.count = epsilon
        self.epsilon = epsilon

    def update(self, batch: NDArray[np.float32]) -> None:
        """Update running statistics with a batch of observations.

        Args:
            batch: Observations array, either (obs_size,) or (batch_size, obs_size)
        """
        batch = np.asarray(batch)
        if batch.ndim == 1:
            batch = batch.reshape(1, -1)
        batch_mean = batch.mean(axis=0)
        batch_var = batch.var(axis=0)
        batch_count = batch.shape[0]
        self._update_from_moments(batch_mean, batch_var, batch_count)

    def _update_from_moments(
        self,
        batch_mean: NDArray[np.float64],
        batch_var: NDArray[np.float64],
        batch_count: int,
    ) -> None:
        """Update statistics using precomputed batch moments."""
        delta = batch_mean - self.mean
        tot_count = self.count + batch_count
        new_mean = self.mean + delta * batch_count / tot_count
        m_a = self.var * self.count
        m_b = batch_var * batch_count
        m_2 = m_a + m_b + np.square(delta) * self.count * batch_count / tot_count
        new_var = m_2 / tot_count
        self.mean = new_mean
        self.var = new_var
        self.count = tot_count

    def normalize(self, x: NDArray[np.float32]) -> NDArray[np.float64]:
        """Normalize observations using running statistics.

        Args:
            x: Observations to normalize

        Returns:
            Normalized observations
        """
        return (x - self.mean) / np.sqrt(self.var + self.epsilon)
