"""
Layer Normalization (LayerNorm) implemented from scratch.
"""

import torch
import torch.nn as nn


class CustomLayerNorm(nn.Module):
    """
    Standard Layer Normalization implemented from PyTorch primitives.
    """

    def __init__(self, dim: int, eps: float = 1e-5, bias: bool = True) -> None:
        """
        Initializes LayerNorm parameters.

        Args:
            dim: Feature hidden dimension.
            eps: Epsilon denominator value for stability.
            bias: Whether to include learnable bias parameter beta.
        """
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
        if bias:
            self.bias = nn.Parameter(torch.zeros(dim))
        else:
            self.register_parameter("bias", None)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass applying mean-centering, variance scaling, and affine transform.
        """
        mean = x.mean(-1, keepdim=True)
        var = x.var(-1, keepdim=True, unbiased=False)
        
        x_norm = (x - mean) / torch.sqrt(var + self.eps)
        
        output = x_norm * self.weight
        if self.bias is not None:
            output = output + self.bias
            
        return output