# SPDX-License-Identifier: Apache-2.0
"""MLX block-scaled quantization backends (affine INT8, MXFP8/4, NVFP4).

Isolated experiment module: probes which ``mx.quantize`` modes the installed
MLX build supports and exposes a thin wrapper around native quantized matmul.
Depends only on ``mlx.core`` and the standard library — do not import the rest
of FastVideo from here.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Final
from collections.abc import Mapping

import mlx.core as mx

# Probe matrix side length: must be divisible by every mode's group size
# (affine g64, mxfp* g32, nvfp4 g16).
_PROBE_DIM: Final[int] = 64


class QuantBackend(str, Enum):
    """Named MLX quantization backends evaluated for M5 Neural Accelerators."""

    AFFINE_INT8_G64 = "affine_int8_g64"
    MXFP8 = "mxfp8"
    MXFP4 = "mxfp4"
    NVFP4 = "nvfp4"


BACKENDS: Final[tuple[str, ...]] = tuple(b.value for b in QuantBackend)

# Backend name -> kwargs for mx.quantize / mx.quantized_matmul.
# Affine baseline matches FastVideo DiT load path (INT8, group size 64).
# MX/NV block-scaled modes use MLX defaults (see mx.quantize docs).
_BACKEND_KWARGS: Final[Mapping[str, Mapping[str, object]]] = {
    QuantBackend.AFFINE_INT8_G64.value: {
        "mode": "affine",
        "bits": 8,
        "group_size": 64,
    },
    QuantBackend.MXFP8.value: {
        "mode": "mxfp8",
        "bits": None,
        "group_size": None,
    },
    QuantBackend.MXFP4.value: {
        "mode": "mxfp4",
        "bits": None,
        "group_size": None,
    },
    QuantBackend.NVFP4.value: {
        "mode": "nvfp4",
        "bits": None,
        "group_size": None,
    },
}

_SUPPORT_CACHE: dict[str, bool] = {}
_SUPPORT_ERROR_CACHE: dict[str, str | None] = {}
_BYTES_CACHE: dict[str, float] = {}


@dataclass(frozen=True)
class QuantizedWeight:
    """Packed quantized weight plus scales/biases for one backend."""

    weight: mx.array
    scales: mx.array
    biases: mx.array | None
    backend: str
    mode: str
    bits: int | None
    group_size: int | None
    # Original (rows, cols) of the fp weight, used for bytes-per-element.
    orig_shape: tuple[int, int]


def _normalize_backend(backend: str) -> str:
    """Normalize a quantization backend name and validate that it is supported.

    Parameters:
        backend (str): Backend name to normalize.

    Returns:
        str: The lowercase backend name without surrounding whitespace.

    Raises:
        ValueError: If the backend name is unknown.
    """
    name = backend.strip().lower()
    if name not in _BACKEND_KWARGS:
        known = ", ".join(BACKENDS)
        raise ValueError(f"Unknown quant backend {backend!r}. Expected one of: {known}")
    return name


def _kwargs_for(backend: str) -> dict[str, object]:
    """Return the MLX quantization arguments configured for a backend.

    Parameters:
        backend (str): Backend name to resolve.

    Returns:
        dict[str, object]: Quantization arguments for the normalized backend.
    """
    return dict(_BACKEND_KWARGS[_normalize_backend(backend)])


def support_error(backend: str) -> str | None:
    """
    Check whether a quantization backend is supported by the current MLX runtime.

    Parameters:
        backend (str): Quantization backend name.

    Returns:
        str | None: An error description when the backend is unsupported, or `None` when supported.
    """
    name = _normalize_backend(backend)
    if name in _SUPPORT_ERROR_CACHE:
        return _SUPPORT_ERROR_CACHE[name]

    kwargs = _kwargs_for(name)
    try:
        w = mx.zeros((_PROBE_DIM, _PROBE_DIM), dtype=mx.float16)
        quantized = mx.quantize(
            w,
            group_size=kwargs["group_size"],  # type: ignore[arg-type]
            bits=kwargs["bits"],  # type: ignore[arg-type]
            mode=str(kwargs["mode"]),
        )
        w_q = quantized[0]
        scales = quantized[1]
        biases = quantized[2] if len(quantized) == 3 else None
        x = mx.zeros((1, _PROBE_DIM), dtype=mx.float16)
        y = mx.quantized_matmul(
            x,
            w_q,
            scales,
            biases,
            transpose=True,
            group_size=kwargs["group_size"],  # type: ignore[arg-type]
            bits=kwargs["bits"],  # type: ignore[arg-type]
            mode=str(kwargs["mode"]),
        )
        mx.eval(y)
        _SUPPORT_ERROR_CACHE[name] = None
        _SUPPORT_CACHE[name] = True
    except Exception as exc:  # noqa: BLE001 - MLX raises varied types per mode/version.
        msg = f"{type(exc).__name__}: {exc}"
        _SUPPORT_ERROR_CACHE[name] = msg
        _SUPPORT_CACHE[name] = False
    return _SUPPORT_ERROR_CACHE[name]


def is_supported(backend: str) -> bool:
    """Return True if the installed MLX build can quantize/matmul with ``backend``."""
    name = _normalize_backend(backend)
    if name not in _SUPPORT_CACHE:
        support_error(name)
    return _SUPPORT_CACHE[name]


def quantize_weight(w: mx.array, backend: str) -> QuantizedWeight:
    """
    Quantize a two-dimensional weight matrix using the specified native MLX backend.

    Parameters:
        w (mx.array): The two-dimensional weight matrix to quantize.
        backend (str): The quantization backend to use.

    Returns:
        QuantizedWeight: The quantized weights and associated quantization metadata.

    Raises:
        ValueError: If the backend is unknown, the weight is not two-dimensional,
            or its last dimension is not divisible by the backend's group size.
        RuntimeError: If the backend is unsupported by the installed MLX build.
    """
    name = _normalize_backend(backend)
    err = support_error(name)
    if err is not None:
        mlx_version = getattr(mx, "__version__", "unknown")
        raise RuntimeError(f"Quant backend {name!r} is not supported by installed mlx "
                           f"({mlx_version}): {err}")

    if w.ndim != 2:
        raise ValueError(f"quantize_weight expects a 2D weight, got shape {tuple(w.shape)}")

    rows, cols = int(w.shape[0]), int(w.shape[1])
    kwargs = _kwargs_for(name)
    group_size = kwargs["group_size"]
    # When group_size is None, MLX applies the mode default; only check when set.
    if isinstance(group_size, int) and cols % group_size != 0:
        raise ValueError(f"Weight last dim {cols} must be divisible by group_size={group_size} "
                         f"for backend {name!r}")

    quantized = mx.quantize(
        w,
        group_size=kwargs["group_size"],  # type: ignore[arg-type]
        bits=kwargs["bits"],  # type: ignore[arg-type]
        mode=str(kwargs["mode"]),
    )
    w_q = quantized[0]
    scales = quantized[1]
    biases = quantized[2] if len(quantized) == 3 else None
    eval_args = [w_q, scales] if biases is None else [w_q, scales, biases]
    mx.eval(*eval_args)

    return QuantizedWeight(
        weight=w_q,
        scales=scales,
        biases=biases,
        backend=name,
        mode=str(kwargs["mode"]),
        bits=kwargs["bits"] if isinstance(kwargs["bits"], int) else None,
        group_size=group_size if isinstance(group_size, int) else None,
        orig_shape=(rows, cols),
    )


def quantized_matmul(x: mx.array, qw: QuantizedWeight) -> mx.array:
    """Compute ``x @ w.T`` in the quantized domain via ``mx.quantized_matmul``."""
    return mx.quantized_matmul(
        x,
        qw.weight,
        qw.scales,
        qw.biases,
        transpose=True,
        group_size=qw.group_size,
        bits=qw.bits,
        mode=qw.mode,
    )


def _artifact_nbytes(qw: QuantizedWeight) -> int:
    """
    Calculate the total storage size of a quantized weight artifact in bytes.

    Parameters:
        qw (QuantizedWeight): Quantized weight artifact whose packed weights, scales, and optional biases are measured.

    Returns:
        int: Total number of bytes used by the artifact's stored arrays.
    """
    total = int(qw.weight.nbytes) + int(qw.scales.nbytes)
    if qw.biases is not None:
        total += int(qw.biases.nbytes)
    return total


def bytes_per_weight(backend: str) -> float:
    """
    Measure the effective storage cost of a quantized weight.

    Parameters:
        backend (str): Quantization backend to measure.

    Returns:
        float: Stored bytes per original weight element, including packed weights,
            scales, and optional biases.

    Raises:
        RuntimeError: If the backend is unsupported.
    """
    name = _normalize_backend(backend)
    if name in _BYTES_CACHE:
        return _BYTES_CACHE[name]

    err = support_error(name)
    if err is not None:
        mlx_version = getattr(mx, "__version__", "unknown")
        raise RuntimeError(f"Cannot measure bytes_per_weight for unsupported backend {name!r} "
                           f"(mlx {mlx_version}): {err}")

    probe = mx.zeros((_PROBE_DIM, _PROBE_DIM), dtype=mx.float16)
    qw = quantize_weight(probe, name)
    n_elem = qw.orig_shape[0] * qw.orig_shape[1]
    value = _artifact_nbytes(qw) / float(n_elem)
    _BYTES_CACHE[name] = value
    return value