"""
Compiler-Optimized Normalization Layers

This module provides normalization layer implementations optimized for PyTorch compiler
and GPU kernel efficiency.

Key Optimizations:
- Fused operations that combine normalization with activation
- Memory-efficient implementations that minimize allocations
- Compiler-friendly patterns that enable automatic optimization
- Numerical stability while maintaining performance
"""


import torch
import torch.nn as nn
import torch.nn.functional as F


class OptimizedLayerNorm(nn.Module):
    """
    Layer normalization optimized for compiler efficiency.

    Improvements over standard LayerNorm:
    - More efficient memory layout
    - Better compiler optimization
    - Optional bias for memory efficiency
    - Stable numerical computation
    """

    def __init__(
        self,
        normalized_shape: int,
        eps: float = 1e-5,
        bias: bool = True,
        device=None,
        dtype=None
    ):
        super().__init__()
        self.normalized_shape = normalized_shape
        self.eps = eps

        self.weight = nn.Parameter(torch.ones(normalized_shape, device=device, dtype=dtype))
        self.bias = nn.Parameter(torch.zeros(normalized_shape, device=device, dtype=dtype)) if bias else None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Optimized layer normalization forward pass.

         GPU OPTIMIZATION DETAILS:
        - Kernel mapping: Uses cuDNN's highly optimized layer normalization kernel
        - Memory access: Single pass through data with fused mean/variance computation
        - Hardware acceleration: Leverages Tensor Cores for mixed precision (fp16/bf16)
        - Vectorization: cuDNN kernel uses vectorized loads/stores for memory efficiency

         WHY THIS OPTIMIZES:
        - F.layer_norm dispatches to cuDNN's hand-optimized assembly kernels
        - Single kernel eliminates intermediate tensor allocations
        - GPU-friendly memory access patterns maximize bandwidth utilization
        - Numerical stability maintained with optimized epsilon handling

        Educational Note: Compare with manual implementation:
        Manual: mean = x.mean(), var = x.var(), norm = (x - mean) / sqrt(var + eps)
        Optimized: Single cuDNN kernel with fused statistics and normalization
        """
        return F.layer_norm(x, (self.normalized_shape,), self.weight, self.bias, self.eps)


class OptimizedRMSNorm(nn.Module):
    """
    RMS Normalization optimized for compiler efficiency.

    RMSNorm is often more efficient than LayerNorm as it doesn't require
    computing the mean, only the root mean square.
    """

    def __init__(
        self,
        normalized_shape: int,
        eps: float = 1e-6,
        device=None,
        dtype=None
    ):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(normalized_shape, device=device, dtype=dtype))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Efficient RMS normalization optimized for modern language models.

         GPU OPTIMIZATION DETAILS:
        - Kernel mapping: Uses PyTorch's optimized norm() kernel + element-wise operations
        - Memory access: Two-pass algorithm but eliminates mean computation overhead
        - Compute efficiency: ~25% fewer operations than LayerNorm (no mean subtraction)
        - Vectorization: All operations are fully vectorizable across feature dimension

         WHY RMS IS FASTER:
        - LayerNorm: compute mean → subtract mean → compute variance → normalize
        - RMSNorm: compute RMS directly → normalize (skips mean computation)
        - Particularly effective for large language models (GPT, LLaMA architectures)
        - Maintains similar training dynamics to LayerNorm with better efficiency

        LayerNorm: (x - mean(x)) / sqrt(var(x) + ε) * γ + β
        RMSNorm:   x / sqrt(mean(x²) + ε) * γ
        """
        #  EDUCATIONAL NOTE: More efficient than LayerNorm - no mean computation needed
        norm = x.norm(dtype=torch.float32, dim=-1, keepdim=True)  # Compute L2 norm efficiently
        rms = norm * (x.shape[-1] ** -0.5)  # Convert to RMS: norm / sqrt(N)
        return (x / (rms + self.eps)).to(x.dtype) * self.weight  # Normalize and scale


class FusedLayerNormActivation(nn.Module):
    """
    Fused Layer Normalization + Activation for maximum efficiency.

    Demonstrates kernel fusion pattern that can provide significant speedups
    by combining operations that would otherwise require separate GPU kernels.
    """

    def __init__(
        self,
        normalized_shape: int,
        activation: str = 'gelu',
        eps: float = 1e-5,
        bias: bool = True
    ):
        super().__init__()
        self.normalized_shape = normalized_shape
        self.eps = eps
        self.activation = activation

        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape)) if bias else None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Fused normalization + activation optimized for kernel fusion.

         GPU OPTIMIZATION DETAILS:
        - Kernel mapping: torch.compile can fuse LayerNorm + Activation into single kernel
        - Memory access: Eliminates intermediate tensor storage between operations
        - Fusion opportunity: Two separate kernel launches → single optimized kernel
        - Hardware utilization: Better memory bandwidth and compute unit usage

         WHY FUSION WORKS:
        - Compiler recognizes producer-consumer pattern (norm → activation)
        - Intermediate results kept in GPU registers instead of global memory
        - Enables more aggressive compiler optimizations (loop fusion, vectorization)
        - Particularly effective with @torch.compile decorator

        Unfused: x → [global memory] → LayerNorm → [global memory] → Activation → output
        Fused:   x → [registers only] → LayerNorm+Activation → output
        """
        #  EDUCATIONAL NOTE: Layer normalization (fusion candidate #1)
        normalized = F.layer_norm(x, (self.normalized_shape,), self.weight, self.bias, self.eps)

        #  EDUCATIONAL NOTE: Activation function (fusion candidate #2)
        # torch.compile will automatically fuse these operations when possible
        if self.activation == 'gelu':
            return F.gelu(normalized)  #  Popular in transformers, complex but very fusable
        elif self.activation == 'relu':
            return F.relu(normalized)  #  Simple, perfect fusion candidate
        elif self.activation == 'swish':
            return normalized * torch.sigmoid(normalized)  #  SiLU activation, good fusion
        else:
            return normalized
