"""
Module that implements a BitNet quantizer for neural networks using JAX/Flax.

This module provides a quantization layer that can be used
to reduce the precision of weights and activations in a neural network,
resulting in more efficient models in terms of memory and computation.

Classes:
    BitNetQuantizer: Implements BitNet quantization.

Dependencies:
    - jax: For array operations and automatic differentiation.
    - flax: For neural network module definitions.
"""

import jax
import jax.numpy as jnp
from flax import linen as nn
import logging

logger = logging.getLogger(__name__)


class BitNetQuantizer(nn.Module):
    """
    A quantization module for BitNet-style quantization.

    This module performs quantization of input arrays to a specified bit width,
    supporting both symmetric and asymmetric quantization, and allowing for
    gradient estimation during backpropagation.

    Attributes:
        bit_width (int): The number of bits to use for quantization.
        symmetric (bool): Whether to use symmetric quantization.
        eps (float): A small value to avoid division by zero.
    """
    bit_width: int = 8
    symmetric: bool = True
    eps: float = 1e-8

    @nn.compact
    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        """
        Perform quantization on the input array.

        Args:
            x (jnp.ndarray): The input array to quantize.

        Returns:
            jnp.ndarray: The quantized array.
        """
        self.validate_bit_width(self.bit_width)

        # Calculate scale and zero_point
        if self.symmetric:
            max_val = jnp.max(jnp.abs(x))
            min_val = -max_val
            zero_point = jnp.zeros_like(max_val)
        else:
            max_val = jnp.max(x)
            min_val = jnp.min(x)
            zero_point = min_val

        scale = (max_val - min_val) / (2**self.bit_width - 1)
        scale = jnp.maximum(scale, self.eps)

        # Quantize
        x_scaled = (x - zero_point) / scale
        x_clipped = jnp.clip(x_scaled, 0, 2**self.bit_width - 1)
        x_rounded = jnp.round(x_clipped)
        x_quantized = x_rounded * scale + zero_point

        # Use straight-through estimator for gradient
        x_quantized = x + jax.lax.stop_gradient(x_quantized - x)

        return x_quantized

    @staticmethod
    def validate_bit_width(bit_width: int):
        """Validate the bit width."""
        if not isinstance(bit_width, int) or bit_width < 2:
            raise ValueError(
                "`bit_width` must be an integer greater than or equal to 2.")

    def get_quantization_params(self, x: jnp.ndarray) -> dict:
        """
        Get the quantization parameters for a given array.

        This method is useful for analysis and debugging.

        Args:
            x (jnp.ndarray): The input array.

        Returns:
            dict: A dictionary containing quantization parameters.
        """
        if self.symmetric:
            max_val = jnp.max(jnp.abs(x))
            min_val = -max_val
            zero_point = jnp.zeros_like(max_val)
        else:
            max_val = jnp.max(x)
            min_val = jnp.min(x)
            zero_point = min_val

        scale = (max_val - min_val) / (2**self.bit_width - 1)
        scale = jnp.maximum(scale, self.eps)

        return {
            "max_val": float(max_val),
            "min_val": float(min_val),
            "scale": float(scale),
            "zero_point": float(zero_point),
            "bit_width": self.bit_width,
            "symmetric": self.symmetric
        }


# Example usage
if __name__ == "__main__":
    # Set up logging
    logging.basicConfig(level=logging.INFO)

    # Create a sample array
    key = jax.random.PRNGKey(0)
    x = jax.random.normal(key, (5, 5))

    # Initialize the quantizer
    quantizer = BitNetQuantizer(bit_width=4, symmetric=True)

    # Perform quantization
    x_quantized = quantizer.apply({}, x)

    # Get quantization parameters
    params = quantizer.get_quantization_params(x)

    print("Original array:")
    print(x)
    print("\nQuantized array:")
    print(x_quantized)
    print("\nQuantization parameters:")
    for key, value in params.items():
        print(f"{key}: {value}")
