import dataclasses
from typing import Literal, NamedTuple

import torch


__all__ = ["LinearCrossEntropyOptions"]


class _AutoDefault(NamedTuple):
    acc_policy: str
    chunking_method: str


# Defaults for the ``"auto"`` sentinel of ``acc_policy`` / ``chunking_method``,
# keyed by ``(device_type, input_dtype)`` -- Pareto picks from an fp64-jacobian
# sweep. CPU picks ``"accurate"`` (only chunked policy with fp32 weight-grad mm
# on CPU; others hit emulated low-precision matmul, ~20-50x slower).
_AUTO_DEFAULTS: dict[tuple[str, torch.dtype], _AutoDefault] = {
    ("cuda", torch.bfloat16): _AutoDefault("compact", "aspect_ratio:2"),
    ("cuda", torch.float16): _AutoDefault("compact", "aspect_ratio:2"),
    ("cpu", torch.bfloat16): _AutoDefault("accurate", "aspect_ratio"),
    ("cpu", torch.float16): _AutoDefault("accurate", "aspect_ratio"),
}
_AUTO_FALLBACK: _AutoDefault = _AutoDefault("compact", "aspect_ratio:2")


@dataclasses.dataclass(slots=True, frozen=True)
class LinearCrossEntropyOptions:
    """Configuration for the chunked implementation of
    :func:`linear_cross_entropy`.

    The chunked implementation processes the batch dimension in pieces, so
    the full ``(num_batches, num_classes)`` logits tensor is never
    materialized -- useful when ``num_classes`` is much larger than
    ``in_features`` (e.g. LLM vocabulary heads). Pass ``options=None`` to
    use the reference path; pass an instance of this class to opt in.

    Zero-argument ``LinearCrossEntropyOptions()`` leaves
    :attr:`acc_policy` and :attr:`chunking_method` set to ``"auto"``,
    resolved at call time from :data:`_AUTO_DEFAULTS` (per-(device, dtype)
    picks measured on A100 / x86 CPU); unlisted pairs fall back to
    ``("compact", "aspect_ratio:2")``.

    Supports a subset of :func:`linear_cross_entropy`; unsupported
    configurations fall through to the reference path with a warning.

    Chunking is a win when ``num_batches >= in_features`` and
    ``num_classes > in_features``; below that, the reference path is
    cheaper.
    """

    allow_retain_graph: bool = False
    """Allow ``retain_graph=True`` on backward. Applies only to the scalar
    reductions (``"mean"`` / ``"sum"``).

    When ``False`` (default), their backward consumes pre-computed gradient
    buffers in place; a second ``.backward()`` raises ``RuntimeError``.

    When ``True``, the buffers are preserved at the cost of one extra
    gradient-sized allocation per call.

    ``reduction="none"`` ignores this field: its backward recomputes the
    chunked gradients from the saved inputs, so ``retain_graph=True`` works
    unconditionally with no extra allocation.

    Higher-order autograd (gradgrad, forward-mode AD) is unsupported.

    Under :func:`torch.compile` this field is auto-promoted to ``True`` for
    the scalar reductions because the default-mode second-backward guard
    relies on a ctx mutation Dynamo doesn't preserve; the wrapper warns on
    the promotion.
    """

    batch_chunk_size: int | None = None
    """Batch rows per chunk. The op loops over
    ``ceil(num_batches / batch_chunk_size)`` chunks; smaller values cut
    peak memory but launch more kernels. Default ``None`` means a single
    chunk. Cannot be combined with :attr:`chunking_method` -- if both are
    set and disagree, ``ValueError`` is raised.
    """

    chunking_method: str | None = "auto"
    """Heuristic for picking :attr:`batch_chunk_size`.

    - ``"auto"`` (default) -- resolves to a per-(device, dtype) pick from
      :data:`_AUTO_DEFAULTS` at call time; falls back to
      ``"aspect_ratio:2"`` for unlisted pairs.
    - ``"aspect_ratio"`` -- sizes each chunk so its
      ``(batch_chunk_size, num_classes)`` logits buffer matches the
      ``(num_batches, in_features)`` input in memory:
      ``next_pow2(ceil(num_batches / ceil(num_classes / in_features)))``.
      Best when ``num_classes >> in_features`` (LLM vocab heads).
    - ``"aspect_ratio:N"`` (``N >= 1``) -- same, divided by ``N``.
      ~N times less peak memory at the cost of N times more chunks.
    - ``None`` -- disables the heuristic; uses :attr:`batch_chunk_size`.
    """

    acc_policy: Literal[
        "accurate",
        "balanced",
        "compact",
        "auto",
    ] = "auto"
    """Precision/memory trade-off for the chunked path. Controls which
    intermediates are kept in :attr:`acc_dtype` vs. the input dtype, and
    whether the per-chunk weight-gradient scratch buffer is materialized.

    - ``"auto"`` (default) -- per-(device, dtype) pick from
      :data:`_AUTO_DEFAULTS`; unlisted pairs fall back to ``"compact"``.
      The fallback assumes a CUDA-like backend with hardware-native
      low-precision matmul; pass ``"accurate"`` explicitly on backends
      that emulate fp16/bf16 GEMMs via fp32 upcast.
    - ``"accurate"`` -- broadest use of :attr:`acc_dtype`; noticeably
      better input-grad accuracy when chunk size is large relative to
      ``num_classes``. Highest peak memory and slowest of the chunked
      policies on CUDA. Only chunked policy whose weight-grad matmul
      runs in fp32 on CPU (other policies hit CPU's emulated
      low-precision path, ~20-50x slower).
    - ``"balanced"`` -- :attr:`acc_dtype` only where needed for gradient
      correctness; keeps a ``(num_classes, in_features)``
      :attr:`acc_dtype` scratch for cross-chunk weight-grad accumulation.
      Same precision as ``"accurate"`` in bf16, slightly looser in fp16,
      faster than ``"accurate"`` in both.
    - ``"compact"`` -- like ``"balanced"`` but drops the weight-grad
      scratch and accumulates per-chunk directly via ``addmm_`` (cuBLAS
      uses an fp32 internal accumulator, so bulk precision matches
      ``"balanced"``). Saves ``num_classes * in_features *
      sizeof(acc_dtype)`` -- typically several hundred MB for an LLM
      head. On non-CUDA mixed-precision falls back to ``"balanced"``.

    Policy effects (``"balanced"`` vs ``"accurate"``) are visible only
    when :attr:`acc_dtype` differs from the input dtype; ``"compact"``
    saves memory in both regimes.
    """

    acc_dtype: torch.dtype | None = None
    """Dtype for internal accumulation. ``None`` resolves at call time
    to ``torch.float32`` under ``acc_policy="auto"`` with fp16/bf16
    input on hardware with mixed-precision mm (CUDA SM 7.0+ for fp16,
    SM 8.0+ for bf16, and CPU); otherwise to the input dtype.
    Mixed-precision currently requires fp16/bf16 input with
    ``acc_dtype=torch.float32``.
    """

    def __post_init__(self):
        if self.acc_policy not in {"auto", "accurate", "balanced", "compact"}:
            raise ValueError(f"invalid acc_policy: {self.acc_policy!r}")
        if self.chunking_method is not None and self.chunking_method != "auto":
            name, sep, factor = self.chunking_method.partition(":")
            if not sep:
                factor = "1"
            if not (name == "aspect_ratio" and factor.isdigit() and int(factor) > 0):
                raise ValueError(f"invalid chunking_method: {self.chunking_method!r}")
        if not (
            self.batch_chunk_size is None
            or (isinstance(self.batch_chunk_size, int) and self.batch_chunk_size > 0)
        ):
            raise ValueError(
                f"batch_chunk_size must be positive int or None, got {self.batch_chunk_size!r}"
            )
        # fp64 is allowed for the internal ``_adjust`` path (fp64 inputs).
        _SUPPORTED_ACC_DTYPES = {
            torch.float16,
            torch.bfloat16,
            torch.float32,
            torch.float64,
        }
        if self.acc_dtype is not None and self.acc_dtype not in _SUPPORTED_ACC_DTYPES:
            raise ValueError(
                f"acc_dtype must be one of {{None, torch.float16, torch.bfloat16, "
                f"torch.float32, torch.float64}}, got {self.acc_dtype!r}. Pass "
                "acc_dtype=None to let the op pick automatically."
            )

    @staticmethod
    def _ceil_div(a: int, b: int) -> int:
        """ceil(a / b) for non-negative integers."""
        return -(-a // b)

    @staticmethod
    def _resolve_auto_acc_dtype(
        dtype: torch.dtype, device: torch.device | None
    ) -> torch.dtype | None:
        """Return fp32 for fp16/bf16 input on hardware with a working
        mixed-precision mm path (CUDA SM 7.0+ for fp16, SM 8.0+ for bf16,
        and CPU); ``None`` otherwise (caller falls back to input dtype).
        """
        if device is None or dtype not in (torch.float16, torch.bfloat16):
            return None
        if device.type == "cuda":
            if not torch.cuda.is_available():
                return None
            major, _ = torch.cuda.get_device_capability(device)
            if dtype == torch.bfloat16 and major < 8:
                return None
            if dtype == torch.float16 and major < 7:
                return None
            return torch.float32
        if device.type == "cpu":
            return torch.float32
        return None

    def _compute_batch_chunk_size(
        self,
        num_batches: int,
        in_features: int,
        num_classes: int,
        method: str | None = None,
    ) -> int:
        """Compute batch_chunk_size from chunking_method given input shapes.

        Pass ``method`` to use a post-``_adjust`` resolved value (e.g. when
        ``self.chunking_method == "auto"``); otherwise uses ``self.chunking_method``.
        To add a method: extend the if/elif chain plus ``__post_init__`` validation.
        """
        method = str(method if method is not None else self.chunking_method)

        if method.startswith("aspect_ratio"):
            factor = int(method.split(":", 1)[1]) if ":" in method else 1
            # See LinearCrossEntropyOptions.chunking_method docstring.
            inc_factor = self._ceil_div(num_classes, in_features)
            target = self._ceil_div(num_batches, inc_factor)
            chunk_size = 1 << (target - 1).bit_length()  # next power of 2 >= target
            return min(chunk_size // factor, num_batches)

        # __post_init__ validates the method, so this is unreachable.
        raise AssertionError(f"unhandled chunking_method: {method!r}")

    def _adjust(self, num_batches, in_features, num_classes, dtype, device=None):
        """Resolve ``"auto"`` sentinels and ``None`` defaults against a
        specific call site; returns a fully concrete options instance.

        Internal API consumed by ``F.linear_cross_entropy``'s chunked
        dispatch and a handful of test call sites. ``device=None``
        forces the :data:`_AUTO_FALLBACK` pick instead of the
        per-device one.
        """
        acc_policy = self.acc_policy
        chunking_method = self.chunking_method
        # Honour an explicit batch_chunk_size by disabling auto chunking
        # (it would conflict with the user's size); acc_policy="auto" is
        # unaffected.
        if chunking_method == "auto" and self.batch_chunk_size is not None:
            chunking_method = None
        if acc_policy == "auto" or chunking_method == "auto":
            if device is not None:
                ap, cm = _AUTO_DEFAULTS.get(
                    (device.type, dtype),
                    _AUTO_FALLBACK,
                )
            else:
                ap, cm = _AUTO_FALLBACK
            if acc_policy == "auto":
                acc_policy = ap
            if chunking_method == "auto":
                chunking_method = cm

        if self.batch_chunk_size is None:
            batch_chunk_size = num_batches
        else:
            batch_chunk_size = min(self.batch_chunk_size, num_batches)

        if chunking_method is not None:
            batch_chunk_size = self._compute_batch_chunk_size(
                num_batches,
                in_features,
                num_classes,
                chunking_method,
            )
            if (
                self.batch_chunk_size is not None
                and self.batch_chunk_size != batch_chunk_size
            ):
                raise ValueError(
                    f"batch_chunk_size (={self.batch_chunk_size}) and "
                    f"chunking_method ('{chunking_method}') give different "
                    f"chunk sizes ({self.batch_chunk_size} vs {batch_chunk_size}); "
                    f"pass only one."
                )

        if self.acc_dtype is None:
            # Under "auto" prefer fp32 on hardware that supports the
            # mixed-precision mm path; otherwise fall back to input dtype.
            auto_acc = (
                self._resolve_auto_acc_dtype(dtype, device)
                if self.acc_policy == "auto"
                else None
            )
            acc_dtype = auto_acc if auto_acc is not None else dtype
        else:
            acc_dtype = self.acc_dtype
        return dataclasses.replace(
            self,
            acc_policy=acc_policy,
            chunking_method=chunking_method,
            batch_chunk_size=max(1, batch_chunk_size),
            acc_dtype=acc_dtype,
        )
