# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
"""LightningModule for RF-DETR training and validation."""

from __future__ import annotations

import importlib
import inspect
import math
import random
import warnings
from typing import Any, Callable, cast

import torch
import torch.nn.functional as F  # noqa: N812 -- project-conventional alias (see AGENTS.md)
from pytorch_lightning import LightningModule, seed_everything
from pytorch_lightning.core.optimizer import LightningOptimizer
from pytorch_lightning.utilities.types import LRSchedulerConfigType, OptimizerLRSchedulerConfig
from torch import Tensor
from torch.optim.lr_scheduler import LRScheduler, ReduceLROnPlateau

from rfdetr._namespace import _namespace_from_configs
from rfdetr.config import (
    ModelConfig,
    TrainConfig,
    _is_managed_optimizer_name,
    _is_managed_scheduler_name,
    _resolve_native_optimizer,
)
from rfdetr.datasets.coco import compute_multi_scale_scales
from rfdetr.models.lwdetr import build_criterion_from_config, build_model_from_config
from rfdetr.models.weights import apply_lora, interpolate_position_embeddings, load_pretrain_weights
from rfdetr.training.callbacks.coco_eval import _get_ema_inner_module
from rfdetr.training.param_groups import get_param_dict
from rfdetr.utilities.logger import get_logger

logger = get_logger()

_OptimizerFactory = Callable[..., torch.optim.Optimizer]

_TRAIN_PROGRESS_LOSS_ALIASES: dict[str, str] = {
    "loss_ce": "loss_cls",
    "loss_bbox": "loss_box",
    "loss_giou": "loss_giou",
    "loss_mask_ce": "mask_ce",
    "loss_mask_dice": "mask_dice",
    "loss_keypoints_l1": "kp_l1",
    "loss_keypoints_findable": "kp_find",
    "loss_keypoints_visible": "kp_vis",
    "loss_keypoints_nll": "kp_nll",
}


def _is_builtin_fused_adamw(optimizer: object) -> bool:
    """Return whether the config selects RF-DETR's built-in (managed, fused) AdamW path.

    Args:
        optimizer: The ``TrainConfig.optimizer`` value (string or callable).

    Returns:
        ``True`` only for the built-in ``"adamw"`` short name.

    Examples:
        >>> _is_builtin_fused_adamw("adamw")
        True
        >>> _is_builtin_fused_adamw("torch.optim.AdamW")
        False
    """
    return isinstance(optimizer, str) and "." not in optimizer and optimizer.strip().lower() == "adamw"


_FUSED_IGNORED_MSG = (
    "fused_optimizer=True is ignored for optimizer=%r; the fused AdamW kernel only applies to the "
    "built-in optimizer='adamw' path."
)


def _import_optimizer_class(dotted_path: str) -> _OptimizerFactory:
    """Import an optimizer class or factory from a dotted path.

    Args:
        dotted_path: Fully-qualified path such as ``"torch.optim.AdamW"`` or
            ``"pytorch_optimizer.Lion"``.

    Returns:
        The imported optimizer class or factory.

    Raises:
        ValueError: If the module or attribute cannot be imported.
    """
    module_path, _, attribute = dotted_path.rpartition(".")
    if not module_path:
        raise ValueError(f"optimizer {dotted_path!r} is not a valid dotted import path.")
    try:
        module = importlib.import_module(module_path)
        return cast(_OptimizerFactory, getattr(module, attribute))
    except (ImportError, AttributeError) as exc:
        raise ValueError(
            f"Could not import optimizer {dotted_path!r}: {exc}. "
            "Use a fully-qualified path to an importable optimizer class, e.g. 'torch.optim.AdamW'."
        ) from exc


def _instantiate_explicit_optimizer(
    optimizer_class: _OptimizerFactory,
    optimizer_name: str,
    param_dicts: list[dict[str, Any]],
    optimizer_kwargs: dict[str, Any],
) -> torch.optim.Optimizer:
    """Instantiate an explicitly-selected optimizer from param groups and kwargs only.

    Explicit optimizers (dotted import paths and callables) receive the RF-DETR
    parameter groups — which already carry per-group learning rates — plus the
    user's ``optimizer_kwargs`` verbatim. RF-DETR injects no ``lr`` or
    ``weight_decay`` of its own.

    Args:
        optimizer_class: Optimizer class or factory to instantiate.
        optimizer_name: Name used in error messages.
        param_dicts: RF-DETR parameter groups with layer-wise learning rates.
        optimizer_kwargs: Keyword arguments forwarded verbatim to the constructor.

    Returns:
        Instantiated optimizer.

    Raises:
        TypeError | ValueError: Re-raised with an RF-DETR-specific hint on failure.
    """
    try:
        return optimizer_class(param_dicts, **optimizer_kwargs)
    except (TypeError, ValueError) as exc:
        raise type(exc)(
            f"Failed to initialize optimizer {optimizer_name!r}: {exc}. "
            "Explicit optimizers (dotted paths and callables) are built from the RF-DETR parameter "
            "groups plus your `optimizer_kwargs` only, with no lr/weight_decay injected. Check that the "
            "class accepts these arguments, or pass a callable/functools.partial needing only `params`."
        ) from exc


def _optimizer_accepts_kwarg(optimizer_class: _OptimizerFactory, name: str) -> bool:
    """Return whether an optimizer constructor accepts a given keyword argument.

    Args:
        optimizer_class: Optimizer class or factory to inspect.
        name: Keyword-argument name to look for.

    Returns:
        ``True`` when the constructor declares ``name`` or accepts arbitrary
        keyword arguments, or when its signature cannot be introspected (e.g. a
        C-implemented constructor) — in which case the constructor validates the
        call itself.

    Examples:
        >>> _optimizer_accepts_kwarg(torch.optim.SGD, "weight_decay")
        True
    """
    try:
        signature = inspect.signature(optimizer_class)
    except (TypeError, ValueError):
        return True
    parameters = signature.parameters.values()
    if any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters):
        return True
    return name in signature.parameters


def _instantiate_optimizer(
    optimizer_class: _OptimizerFactory,
    optimizer_name: str,
    param_dicts: list[dict[str, Any]],
    train_config: TrainConfig,
) -> torch.optim.Optimizer:
    """Instantiate an optimizer class with RF-DETR optimizer arguments.

    ``weight_decay`` is injected only when the optimizer constructor accepts it,
    so optimizers with a different regularization API are not forced to fail.

    Args:
        optimizer_class: Optimizer class or factory to instantiate.
        optimizer_name: Name used in error messages.
        param_dicts: RF-DETR parameter groups with layer-wise learning rates.
        train_config: Training config with base optimizer hyperparameters.

    Returns:
        Instantiated optimizer.

    Raises:
        TypeError | ValueError: Re-raised with an RF-DETR-specific hint when the
            optimizer constructor rejects the supplied arguments.
    """
    init_kwargs: dict[str, Any] = {"lr": train_config.lr}
    if _optimizer_accepts_kwarg(optimizer_class, "weight_decay"):
        init_kwargs["weight_decay"] = train_config.weight_decay
    init_kwargs.update(train_config.optimizer_kwargs)
    try:
        return optimizer_class(param_dicts, **init_kwargs)
    except (TypeError, ValueError) as exc:
        raise type(exc)(
            f"Failed to initialize optimizer {optimizer_name!r}: {exc}. "
            "For managed torch.optim short names, RF-DETR passes `params`, `lr`, and (when supported) "
            "`weight_decay`, then your `optimizer_kwargs`; this usually means an unsupported entry in "
            "`optimizer_kwargs`."
        ) from exc


_SchedulerFactory = Callable[..., LRScheduler | ReduceLROnPlateau]


def _import_scheduler_class(dotted_path: str) -> _SchedulerFactory:
    """Import an LR-scheduler class or factory from a dotted path.

    Args:
        dotted_path: Fully-qualified path such as ``"torch.optim.lr_scheduler.StepLR"`` or
            ``"pytorch_optimizer.CosineAnnealingWarmupRestarts"``.

    Returns:
        The imported scheduler class or factory.

    Raises:
        ValueError: If the module or attribute cannot be imported.
    """
    module_path, _, attribute = dotted_path.rpartition(".")
    if not module_path:
        raise ValueError(f"lr_scheduler {dotted_path!r} is not a valid dotted import path.")
    try:
        module = importlib.import_module(module_path)
        return cast(_SchedulerFactory, getattr(module, attribute))
    except (ImportError, AttributeError) as exc:
        raise ValueError(
            f"Could not import lr_scheduler {dotted_path!r}: {exc}. "
            "Use a fully-qualified path to an importable scheduler class, e.g. 'torch.optim.lr_scheduler.StepLR'."
        ) from exc


def _instantiate_explicit_scheduler(
    scheduler_factory: _SchedulerFactory,
    scheduler_name: str,
    optimizer: torch.optim.Optimizer,
    scheduler_kwargs: dict[str, Any],
) -> LRScheduler | ReduceLROnPlateau:
    """Instantiate an explicitly-selected LR scheduler from the optimizer and kwargs only.

    Explicit schedulers (dotted import paths and callables) receive the built optimizer plus the
    user's ``lr_scheduler_kwargs`` verbatim. RF-DETR injects no ``total_steps`` / ``T_max`` of its
    own — the managed ``"step"`` / ``"cosine"`` presets remain the runtime-aware option.

    Args:
        scheduler_factory: Scheduler class or factory to instantiate.
        scheduler_name: Name used in error messages.
        optimizer: The optimizer the scheduler drives.
        scheduler_kwargs: Keyword arguments forwarded verbatim to the constructor.

    Returns:
        Instantiated scheduler.

    Raises:
        TypeError | ValueError: Re-raised with an RF-DETR-specific hint on failure.
    """
    try:
        return scheduler_factory(optimizer, **scheduler_kwargs)
    except (TypeError, ValueError) as exc:
        raise type(exc)(
            f"Failed to initialize lr_scheduler {scheduler_name!r}: {exc}. "
            "Explicit schedulers (dotted paths and callables) are built from the optimizer plus your "
            "`lr_scheduler_kwargs` only, with no total_steps/T_max injected. Check that the class accepts these "
            "arguments, or pass a callable/functools.partial needing only the optimizer."
        ) from exc


def _build_managed_scheduler(
    optimizer: torch.optim.Optimizer,
    train_config: TrainConfig,
    total_steps: int,
    steps_per_epoch: int,
    warmup_steps: int,
) -> torch.optim.lr_scheduler.LambdaLR:
    """Build the managed ``"step"`` / ``"cosine"`` scheduler as a warmup-aware ``LambdaLR``.

    Preserves RF-DETR's built-in schedule: a linear warmup ramp over ``warmup_steps`` followed by
    either cosine annealing down to ``min_factor`` or a 10x step decay after ``lr_drop`` epochs. The
    ``min_factor`` and ``lr_drop`` values are read from ``lr_scheduler_kwargs`` first (the current API),
    falling back to the deprecated ``lr_min_factor`` / ``lr_drop`` fields.

    Args:
        optimizer: The optimizer the scheduler drives.
        train_config: Training config carrying the preset name and schedule knobs.
        total_steps: Total optimizer steps over the whole run.
        steps_per_epoch: Optimizer steps per epoch.
        warmup_steps: Number of optimizer steps in the linear warmup ramp.

    Returns:
        A ``LambdaLR`` implementing the managed schedule.
    """
    kwargs = train_config.lr_scheduler_kwargs
    # Managed presets are always strings (guaranteed by the _is_managed_scheduler_name branch at the call site).
    preset = cast(str, train_config.lr_scheduler).strip().lower()
    min_factor = float(kwargs.get("min_factor", train_config.lr_min_factor))
    lr_drop = int(kwargs.get("lr_drop", train_config.lr_drop))

    def lr_lambda(current_step: int) -> float:
        if current_step < warmup_steps:
            return float(current_step) / float(max(1, warmup_steps))
        if preset == "cosine":
            progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
            return min_factor + (1 - min_factor) * 0.5 * (1 + math.cos(math.pi * progress))
        # Step decay: drop by 10x after lr_drop epochs.
        if current_step < lr_drop * steps_per_epoch:
            return 1.0
        return 0.1

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)


def _wrap_with_warmup(
    scheduler: torch.optim.lr_scheduler.LRScheduler,
    optimizer: torch.optim.Optimizer,
    warmup_steps: int,
) -> torch.optim.lr_scheduler.SequentialLR:
    """Prepend a linear warmup ramp to an explicit scheduler via ``SequentialLR``.

    The warmup ramps the LR from ``1 / warmup_steps`` of its base value up to full over
    ``warmup_steps`` optimizer steps, then hands control to ``scheduler``. ``ReduceLROnPlateau`` is
    metric-driven and cannot be composed this way — callers must skip wrapping it.

    Args:
        scheduler: The explicit scheduler to run after warmup.
        optimizer: The optimizer both schedulers drive.
        warmup_steps: Number of optimizer steps in the linear warmup ramp.

    Returns:
        A ``SequentialLR`` chaining the warmup ramp and ``scheduler``.
    """
    warmup = torch.optim.lr_scheduler.LinearLR(
        optimizer,
        start_factor=1.0 / max(1, warmup_steps),
        end_factor=1.0,
        total_iters=warmup_steps,
    )
    return torch.optim.lr_scheduler.SequentialLR(
        optimizer,
        schedulers=[warmup, scheduler],
        milestones=[warmup_steps],
    )


class RFDETRModelModule(LightningModule):
    """LightningModule wrapping the RF-DETR model and training loop.

    Args:
        model_config: Architecture configuration.
        train_config: Training hyperparameter configuration.
    """

    def __init__(self, model_config: ModelConfig, train_config: TrainConfig) -> None:
        super().__init__()
        self.model_config = model_config
        self.train_config = train_config
        # Manual optimization is enabled only for keypoint models so that the box-count
        # normalizer can be accumulated across grad-accum microbatches. Detection and
        # segmentation use Lightning's automatic optimization (PTL handles accumulation,
        # AMP, and gradient clipping), which keeps their step semantics unchanged from
        # the pre-fix/scaling behaviour.
        self._use_manual_optimization: bool = bool(getattr(model_config, "use_grouppose_keypoints", False))
        self.automatic_optimization = not self._use_manual_optimization
        # LR-scheduler stepping cadence resolved in configure_optimizers(); read by the manual-optimization
        # step loop and the epoch-end hook. Defaults keep pre-configure behaviour (per-step stepping).
        self._lr_scheduler_interval: str = "step"
        self._lr_scheduler_monitor: str | None = None
        self._accumulated_box_normalizer: Tensor | None = None
        # Allow partial state-dict loading when resuming from a .pth checkpoint
        # (which contains only model weights, not criterion/postprocess state).
        self.strict_loading = False

        # Model, criterion, and postprocessor.
        self.model = build_model_from_config(model_config, train_config)
        if model_config.pretrain_weights is not None:
            # Canonical loader handles PE interpolation, PTL .ckpt normalisation,
            # per-group query slicing, class-name extraction, partial-load warnings,
            # and writes any auto-aligned ``num_classes`` back onto ``model_config``.
            load_pretrain_weights(self.model, self.model_config)
            if model_config.use_grouppose_keypoints:
                # Older model shims may omit the keypoint reset hook; call it only when implemented.
                reset_keypoint_gaussian_parameters = getattr(self.model, "reset_keypoint_gaussian_parameters", None)
                if callable(reset_keypoint_gaussian_parameters):
                    reset_keypoint_gaussian_parameters()
                    logger.info(
                        "Reset keypoint Gaussian precision outputs to unit values after pretrained weight load."
                    )
        if model_config.backbone_lora:
            apply_lora(self.model)

        # Build criterion/postprocessors after potential num_classes alignment so
        # they are constructed with a config that matches the current model head.
        self.criterion, self.postprocess = build_criterion_from_config(self.model_config, self.train_config)

        # torch.compile is opt-in: set model_config.compile=True to enable.
        # Only enabled on CUDA; MPS and CPU do not benefit from compilation.
        # Use the fork-safe DEVICE constant instead of torch.cuda.is_available(),
        # which creates a CUDA driver context that breaks fork-based DDP.
        from rfdetr.config import DEVICE

        accelerator = str(train_config.accelerator).lower()
        uses_cuda_accelerator = accelerator in {"auto", "gpu", "cuda"}
        compile_enabled = (
            model_config.compile and DEVICE == "cuda" and uses_cuda_accelerator and not train_config.multi_scale
        )
        if model_config.compile and train_config.multi_scale:
            logger.info(
                "Disabling torch.compile: multi_scale=True causes dynamic shapes "
                "(incompatible with XLA -- each scale = separate graph trace). "
                "Use do_random_resize_via_padding=True on TPU to avoid recompilation overhead."
            )
        if compile_enabled:
            # dynamic=True: one compiled graph handles all multi-scale input sizes instead
            # of recompiling per (H, W) pair. suppress_errors=True: if inductor can't
            # compile a subgraph (e.g. bicubic backward with symbolic shapes), it falls
            # back to eager mode for that subgraph rather than crashing.
            # capture_scalar_outputs=True: include Tensor.item() calls
            # (gen_encoder_output_proposals / ms_deform_attn use spatial-shape .item()
            # as Python slice indices). Safe with dynamic=True because item() results
            # are backed symbols derived from input shapes — not unbacked symbols that
            # would cause PendingUnbackedSymbolNotFound (which only occurs without dynamic).
            torch._dynamo.config.suppress_errors = True
            torch._dynamo.config.capture_scalar_outputs = True
            # OptimizedModule forwards attribute access to the wrapped LWDETR via
            # __getattr__ at runtime, so self.model keeps working everywhere it's used below.
            self.model = torch.compile(self.model, dynamic=True)  # type: ignore[assignment]

    # ------------------------------------------------------------------
    # PTL lifecycle hooks
    # ------------------------------------------------------------------

    def on_fit_start(self) -> None:
        """Seed RNGs at fit start when ``TrainConfig.seed`` is set.

        This avoids hidden global side-effects in ``build_trainer`` while still preserving deterministic training
        behaviour for actual fit runs.
        """
        if self.train_config.seed is not None:
            seed_everything(self.train_config.seed + self.global_rank, workers=True)

    def on_train_start(self) -> None:
        """Normalize restored fused-optimizer state before the first training step.

        Lightning restores optimizer state after ``on_fit_start``.  Fused AdamW is strict about the dtype, device, and
        layout of its moment tensors, so resuming from a checkpoint can fail if Lightning rehydrates those tensors in a
        layout that no longer matches the live parameters.  Recasting same-shaped floating-point tensors here keeps the
        resumed optimizer compatible without discarding the saved momentum state.
        """
        if not self._use_fused_optimizer:
            return

        try:
            optimizers = self.optimizers(use_pl_optimizer=False)
        except RuntimeError:
            return
        if optimizers is None:
            return
        if isinstance(optimizers, list):
            optimizer_list = optimizers
        else:
            optimizer_list = [optimizers]

        normalized_tensors = 0
        for optimizer in optimizer_list:
            if not isinstance(optimizer, torch.optim.Optimizer):
                optimizer = getattr(optimizer, "optimizer", optimizer)
            if isinstance(optimizer, torch.optim.Optimizer):
                normalized_tensors += self._normalize_optimizer_state(optimizer)
        if normalized_tensors and getattr(self.trainer, "is_global_zero", True):
            logger.info(
                "Normalized %d restored fused AdamW state tensors after checkpoint resume.",
                normalized_tensors,
            )

    def on_train_batch_start(self, batch: tuple[Any, Any], batch_idx: int) -> None:
        """Apply optional multi-scale resize to the incoming batch.

        Modifications to ``batch`` (in-place on ``NestedTensor``) are visible in ``training_step`` because they share
        the same object.

        Args:
            batch: Tuple of (NestedTensor samples, list of target dicts).
            batch_idx: Index of the current batch within the epoch.
        """
        tc = self.train_config
        mc = self.model_config

        if tc.multi_scale and not tc.do_random_resize_via_padding:
            samples, _ = batch
            scales = compute_multi_scale_scales(mc.resolution, tc.expanded_scales, mc.patch_size, mc.num_windows)
            step = self.trainer.global_step
            # Use a step-local generator so the scale choice is deterministic and DDP-consistent
            # without reseeding the process-global RNG on every batch.
            scale = random.Random(step).choice(scales)
            with torch.no_grad():
                samples.tensors = F.interpolate(samples.tensors, size=scale, mode="bilinear", align_corners=False)
                samples.mask = (
                    F.interpolate(samples.mask.unsqueeze(1).float(), size=scale, mode="nearest").squeeze(1).bool()
                )

    def on_train_epoch_start(self) -> None:
        """Reset the accumulated box normalizer at the start of every training epoch.

        Lightning may reuse the module across epochs without calling ``_step_optimizer`` at the boundary (for example
        when an epoch ends mid-accumulation window with a non-divisible batch count). Clearing the accumulator here
        guarantees the manual-optimization path always starts each epoch from a known state, so the first microbatch's
        gradients are scaled by its own box count and not by a stale previous-epoch denominator.

        This is a no-op for non-keypoint models because they use Lightning's automatic optimization path and never
        populate ``self._accumulated_box_normalizer``.

        Note: on finite datasets the final-batch fallback in ``_should_step_optimizer`` always flushes a partial
        trailing window, so this reset is the only change needed.  On IterableDatasets (infinite
        ``num_training_batches``) a partial window may survive epoch end with un-stepped gradients; those are
        discarded here and the optimizer is zeroed so the first microbatch of the new epoch starts from a clean state.
        """
        if self._accumulated_box_normalizer is not None:
            # Discard any partial accumulation window that survived the epoch boundary
            # (only possible for IterableDatasets where num_training_batches is infinite).
            try:
                opts = self.optimizers()
                for opt in opts if isinstance(opts, list) else [opts]:
                    opt.zero_grad()
            except RuntimeError:
                pass  # Not attached to Trainer (unit-test context); nothing to zero.
        self._accumulated_box_normalizer = None

    def training_step(self, batch: tuple[Any, Any], batch_idx: int) -> Tensor | dict[str, Any]:
        """Compute loss for one training step and log metrics.

        PTL handles AMP (``precision``) without a manual ``GradScaler``. Keypoint models perform manual optimization so
        box-count loss normalization is based on the full accumulated effective batch rather than each microbatch
        independently; detection and segmentation models keep Lightning's automatic optimization path.

        Args:
            batch: Tuple of (NestedTensor samples, list of target dicts).
            batch_idx: Batch index within the epoch.

        Returns:
            Scalar loss tensor by default. When ``compute_train_metrics=True``,
            returns a Lightning-compatible dict containing ``loss`` plus
            detached postprocessed predictions for train mAP logging.
        """
        samples, targets = batch
        batch_size = len(targets)
        outputs = self.model(samples, targets)
        if self._use_manual_optimization:
            loss_dict, raw_loss, normalizer = self._compute_train_losses(outputs, targets)
            loss_for_backward = self._scale_loss_for_accumulation(raw_loss, normalizer)
        else:
            loss_dict = self.criterion(outputs, targets)
            loss_for_backward = None
        weight_dict = self.criterion.weight_dict
        loss: Tensor = torch.stack([loss_dict[k] * weight_dict[k] for k in loss_dict if k in weight_dict]).sum()
        # Automatic optimization path: divide by accumulate_grad_batches so the accumulated
        # gradient matches a single large batch, matching the legacy engine.  PTL accumulates
        # full-scale gradients by default; dividing here keeps the effective LR identical.
        accumulate_grad_batches = max(1, int(self.trainer.accumulate_grad_batches))
        loss_for_return = loss if self._use_manual_optimization else loss / accumulate_grad_batches
        train_log_sync_dist = bool(self.train_config.train_log_sync_dist)
        train_log_on_step = bool(self.train_config.train_log_on_step)
        self.log_dict(
            {f"train/{k}": v for k, v in loss_dict.items()},
            on_step=train_log_on_step,
            on_epoch=True,
            sync_dist=train_log_sync_dist,
            batch_size=batch_size,
        )
        self.log(
            "train/loss",
            loss,
            prog_bar=True,
            on_step=train_log_on_step,
            on_epoch=True,
            sync_dist=train_log_sync_dist,
            batch_size=batch_size,
        )
        self._log_train_progress_metrics(loss, loss_dict, batch_size=batch_size)
        optimizer = self.optimizers()
        if isinstance(optimizer, list):
            optimizer = optimizer[0]
        # Optimizer may have multiple param groups with different LRs (e.g., backbone/decoder).
        # Preserve the first group's LR for backward compatibility, but also log the
        # min/max across all groups so the progress bar reflects the full schedule.
        group_lrs = [pg["lr"] for pg in optimizer.param_groups if "lr" in pg]
        if group_lrs:
            base_lr = group_lrs[0]
            min_lr = min(group_lrs)
            max_lr = max(group_lrs)
            self.log("train/lr", base_lr, prog_bar=False, on_step=True, on_epoch=False)
            self.log("train/lr_min", min_lr, prog_bar=False, on_step=True, on_epoch=False)
            self.log("train/lr_max", max_lr, prog_bar=False, on_step=True, on_epoch=False)
        if self._use_manual_optimization:
            # loss_for_backward is only None in the automatic-optimization branch above,
            # which is mutually exclusive with _use_manual_optimization.
            assert loss_for_backward is not None
            self.manual_backward(loss_for_backward)
            if self._should_step_optimizer(batch_idx):
                self._step_optimizer(optimizer)
        if self.train_config.compute_train_metrics:
            with torch.no_grad():
                orig_sizes = torch.stack([t["orig_size"] for t in targets])
                # Slice to group-0 queries only — mirrors the eval-mode path in
                # lwdetr.py that trims refpoint_embed to [:num_queries]. Without
                # this, training mode emits group_detr×num_queries queries (e.g.
                # 13×300=3900) and postprocess top-k selection draws from all
                # groups, producing OKS/mAP values ~50× below true accuracy.
                nq = self.model_config.num_queries
                # Only include tensor-valued keys — pred_masks is a dict in
                # train mode (sparse_forward) and postprocess cannot handle it.
                inference_outputs = {
                    k: v[:, :nq] if v.ndim >= 2 else v
                    for k, v in outputs.items()
                    if k in ("pred_logits", "pred_boxes", "pred_masks", "pred_keypoints") and isinstance(v, Tensor)
                }
                results = self.postprocess(inference_outputs, orig_sizes)
            return {
                "loss": loss_for_return.detach() if self._use_manual_optimization else loss_for_return,
                "results": self._detach_results(results),
                "targets": targets,
            }
        return loss_for_return.detach() if self._use_manual_optimization else loss_for_return

    def _compute_train_losses(
        self,
        outputs: dict[str, Tensor],
        targets: list[dict[str, Tensor]],
    ) -> tuple[dict[str, Tensor], Tensor, Tensor]:
        """Compute normalized losses for logging and raw weighted loss for backward.

        Args:
            outputs: Model output dictionary.
            targets: Target dictionaries for the current batch.

        Returns:
            A tuple of normalized loss dictionary, unnormalized weighted loss numerator, and box normalizer.
        """
        weight_dict = self.criterion.weight_dict
        if not getattr(self.criterion, "supports_loss_normalizer_override", False):
            raise ValueError(
                f"{type(self.criterion).__name__}.supports_loss_normalizer_override is False; "
                "manual optimization (keypoint models) requires a criterion that accepts a "
                "num_boxes keyword argument. Set supports_loss_normalizer_override = True on "
                "your criterion subclass and implement the num_boxes parameter in forward()."
            )
        normalizer = self.criterion.num_boxes_for_targets(outputs, targets)
        numerator_loss_dict = self.criterion(outputs, targets, num_boxes=torch.ones_like(normalizer))
        # Keys in weight_dict are loss terms whose criterion implementation divides by num_boxes
        # (so passing num_boxes=1.0 yields raw numerators that we divide by normalizer here).
        # Keys outside weight_dict (e.g. "class_error", "cardinality_error") are diagnostics
        # that do NOT divide by num_boxes internally — they are passed through unchanged.
        # If a future loss term divides by num_boxes AND is omitted from weight_dict, its
        # logged value will be on a different scale than the keypoint path; verify when adding
        # new criterion terms.
        loss_dict = {
            key: value / normalizer if key in weight_dict else value for key, value in numerator_loss_dict.items()
        }
        raw_loss = sum(numerator_loss_dict[k] * weight_dict[k] for k in numerator_loss_dict if k in weight_dict)
        return loss_dict, raw_loss, normalizer

    def _scale_loss_for_accumulation(
        self,
        raw_loss: Tensor,
        normalizer: Tensor,
    ) -> Tensor:
        """Scale the current numerator loss by the accumulated box denominator.

        Args:
            raw_loss: Current microbatch weighted loss numerator.
            normalizer: Current microbatch box denominator.

        Returns:
            Loss scalar to pass to ``manual_backward``.
        """
        normalizer = normalizer.detach()
        previous_normalizer = self._accumulated_box_normalizer
        accumulated_normalizer = normalizer if previous_normalizer is None else previous_normalizer + normalizer
        if previous_normalizer is not None:
            self._rescale_accumulated_gradients(previous_normalizer / accumulated_normalizer)
        self._accumulated_box_normalizer = accumulated_normalizer.detach()
        return raw_loss / accumulated_normalizer

    def _rescale_accumulated_gradients(self, scale: Tensor) -> None:
        """Rescale gradients already accumulated in the current optimizer window.

        Args:
            scale: Multiplicative factor that converts previous gradients from the old denominator to the new one.
        """
        for parameter in self.parameters():
            if parameter.grad is not None:
                parameter.grad.mul_(scale.to(device=parameter.grad.device, dtype=parameter.grad.dtype))

    def _should_step_optimizer(self, batch_idx: int) -> bool:
        """Return whether the current batch closes an optimizer accumulation window.

        The optimizer steps when either:

        - The current batch closes a complete ``grad_accum_steps`` window
          (``(batch_idx + 1) % grad_accum_steps == 0``), or
        - This is the final batch of the epoch and a partial accumulation window
          is still open, so the trailing microbatches are not silently dropped.

        Lightning's ``Trainer.num_training_batches`` may be reported as ``float('inf')``
        for iterable / streaming datasets where the epoch length is unknown. In that case
        only the modulo path can ever close the window — the final-batch fallback is
        skipped because ``batch_idx + 1`` can never reach infinity.

        Args:
            batch_idx: Batch index within the epoch.

        Returns:
            ``True`` when the optimizer should step after this batch.
        """
        accum_steps = max(1, int(self.train_config.grad_accum_steps))
        if (batch_idx + 1) % accum_steps == 0:
            return True
        num_training_batches = getattr(self.trainer, "num_training_batches", None)
        return (
            isinstance(num_training_batches, (int, float))
            and math.isfinite(num_training_batches)
            and batch_idx + 1 >= num_training_batches
        )

    def _step_optimizer(self, optimizer: torch.optim.Optimizer | LightningOptimizer) -> None:
        """Clip gradients, step optimizer and scheduler, then reset accumulation state.

        Args:
            optimizer: Optimizer returned by Lightning.
        """
        trainer_gradient_clip_val = getattr(self.trainer, "gradient_clip_val", None)
        if trainer_gradient_clip_val is None:
            gradient_clip_val = self.train_config.clip_max_norm
        elif isinstance(trainer_gradient_clip_val, (int, float)):
            gradient_clip_val = trainer_gradient_clip_val
        else:
            gradient_clip_val = None
        gradient_clip_algorithm = getattr(self.trainer, "gradient_clip_algorithm", None)
        if not isinstance(gradient_clip_algorithm, str):
            gradient_clip_algorithm = None
        if gradient_clip_val is not None and gradient_clip_val > 0:
            self.clip_gradients(
                optimizer,
                gradient_clip_val=gradient_clip_val,
                gradient_clip_algorithm=gradient_clip_algorithm,
            )
        optimizer.step()
        optimizer.zero_grad()
        self._step_lr_scheduler()
        self._accumulated_box_normalizer = None

    def _current_lr_scheduler(self) -> LRScheduler | ReduceLROnPlateau | None:
        """Return the single configured LR scheduler, or ``None`` when none is available.

        Returns:
            The scheduler object (unwrapping Lightning's single-element list), or ``None`` when no
            scheduler is configured yet or Lightning is between fit stages.
        """
        try:
            scheduler = self.lr_schedulers()
        except (AttributeError, RuntimeError):
            return None
        if isinstance(scheduler, list):
            return scheduler[0] if scheduler else None
        return scheduler

    def _step_lr_scheduler(self) -> None:
        """Step step-interval schedulers once per optimizer step (manual-optimization path).

        Epoch-interval schedulers and metric-driven ``ReduceLROnPlateau`` are stepped at epoch boundaries by
        ``on_train_epoch_end`` / ``on_validation_epoch_end`` instead, so they are skipped here.
        """
        if self._lr_scheduler_interval != "step":
            return
        scheduler = self._current_lr_scheduler()
        if scheduler is None or isinstance(scheduler, ReduceLROnPlateau):
            return
        scheduler.step()

    def on_train_epoch_end(self) -> None:
        """Step epoch-interval (non-plateau) schedulers on the manual-optimization path.

        The automatic-optimization path leaves scheduler stepping entirely to Lightning; only the manual keypoint loop
        steps schedulers itself.
        """
        if self.automatic_optimization or self._lr_scheduler_interval != "epoch":
            return
        scheduler = self._current_lr_scheduler()
        if scheduler is None or isinstance(scheduler, ReduceLROnPlateau):
            return
        scheduler.step()

    def on_validation_epoch_end(self) -> None:
        """Step ``ReduceLROnPlateau`` from the monitored metric on the manual-optimization path.

        The automatic-optimization path lets Lightning feed the monitored metric; the manual keypoint loop must read it
        from ``trainer.callback_metrics`` and step the scheduler itself. The pre-training sanity-check validation is
        skipped so plateau patience/cooldown bookkeeping is not seeded from the untrained model.
        """
        if self.automatic_optimization or self.trainer.sanity_checking:
            return
        scheduler = self._current_lr_scheduler()
        if not isinstance(scheduler, ReduceLROnPlateau):
            return
        monitor = self._lr_scheduler_monitor or "val/loss"
        metric = self.trainer.callback_metrics.get(monitor)
        if metric is None:
            # Warn-and-continue would let the LR never reduce while training silently proceeds. Fail loud instead,
            # mirroring Lightning's strict-monitor behavior on the automatic-optimization path.
            raise RuntimeError(
                f"ReduceLROnPlateau monitor {monitor!r} was not found in callback_metrics, so the learning rate "
                "would never be reduced. Ensure the monitored metric is logged every validation epoch (e.g. set "
                "compute_val_loss=True for the default 'val/loss' monitor), or set lr_scheduler_monitor to a metric "
                "that is produced."
            )
        scheduler.step(metric)

    @staticmethod
    def _detach_results(results: list[dict[str, Tensor]]) -> list[dict[str, Tensor]]:
        """Detach postprocessed result tensors before handing them to callbacks.

        Args:
            results: Per-image postprocessed prediction dictionaries.

        Returns:
            Per-image dictionaries with tensor values detached from the graph.
        """
        return [
            {key: value.detach() if torch.is_tensor(value) else value for key, value in result.items()}
            for result in results
        ]

    def _log_train_progress_metrics(
        self,
        loss: Tensor,
        loss_dict: dict[str, Tensor],
        *,
        batch_size: int,
    ) -> None:
        """Log compact per-step convergence metrics for the progress bar only.

        Args:
            loss: Unscaled aggregate training loss.
            loss_dict: Raw criterion loss dictionary.
            batch_size: Current batch size used by Lightning for metric reduction metadata.
        """
        # When ``train_log_on_step`` is True, the ``train/loss`` call in ``training_step``
        # logs with ``on_step=True, on_epoch=True``; Lightning forks that into
        # ``train/loss_step`` + ``train/loss_epoch``, and ``train/loss_step`` already
        # provides the live per-step progress-bar view. Emitting this separate ``loss``
        # scalar in that case just duplicates it, so only log it on the default
        # ``train_log_on_step=False`` path.
        if not bool(self.train_config.train_log_on_step):
            self.log(
                "loss",
                loss,
                prog_bar=True,
                logger=False,
                on_step=True,
                on_epoch=False,
                batch_size=batch_size,
            )
        for loss_name, progress_name in _TRAIN_PROGRESS_LOSS_ALIASES.items():
            value = loss_dict.get(loss_name)
            if value is None:
                continue
            self.log(
                progress_name,
                value,
                prog_bar=True,
                logger=False,
                on_step=True,
                on_epoch=False,
                batch_size=batch_size,
            )

    def _log_val_loss_metrics(
        self,
        loss: Tensor,
        loss_dict: dict[str, Tensor],
        *,
        batch_size: int,
    ) -> None:
        """Log aggregate and component validation losses.

        Args:
            loss: Aggregate weighted validation loss.
            loss_dict: Raw criterion loss dictionary.
            batch_size: Current batch size used by Lightning for metric reduction metadata.
        """
        self.log_dict(
            {f"val/{k}": v for k, v in loss_dict.items()},
            on_step=False,
            on_epoch=True,
            sync_dist=True,
            batch_size=batch_size,
        )
        self.log("val/loss", loss, prog_bar=True, on_epoch=True, sync_dist=True, batch_size=batch_size)

    def _resolve_eval_model(self) -> Any:
        """Return the model to forward through for validation.

        When ``TrainConfig.eval_ema_only`` is set, validation forwards through the EMA-averaged
        weights directly instead of the base model — this replaces the second, duplicate
        base+EMA forward pass ``COCOEvalCallback`` would otherwise also run every validation
        batch (see issue 416). Falls back to the base model when EMA is enabled but not yet
        warmed up (e.g. the very first validation epoch, before ``RFDETREMACallback.setup``
        has built its averaged model), or when this module isn't attached to a ``Trainer`` at
        all (``LightningModule.trainer`` raises ``RuntimeError`` rather than returning ``None``
        when unattached — e.g. ``validation_step`` called directly outside ``Trainer.fit``/
        ``Trainer.validate``).

        Uses the same ``_get_ema_inner_module`` helper as ``COCOEvalCallback`` (see
        ``coco_eval.py``) so both consumers resolve the EMA-averaged detection net through one
        shared code path instead of independently duck-typing ``RFDETREMACallback``.

        Returns:
            The base model, or the EMA-averaged inner module's underlying detection net when
            ``eval_ema_only`` is active and available.
        """
        if not self.train_config.eval_ema_only:
            return self.model
        try:
            callbacks = getattr(self.trainer, "callbacks", [])
        except RuntimeError:
            return self.model
        for callback in callbacks:
            ema_inner = _get_ema_inner_module(callback)
            if ema_inner is not None:
                return ema_inner.model
        return self.model

    def validation_step(self, batch: tuple[Any, Any], batch_idx: int) -> dict[str, Any]:
        """Run forward pass and postprocess for one validation step.

        Returns raw results and targets so ``COCOEvalCallback`` can accumulate them across the epoch via
        ``on_validation_batch_end``.

        Args:
            batch: Tuple of (NestedTensor samples, list of target dicts).
            batch_idx: Batch index within the validation epoch.

        Returns:
            Dict with ``results`` (postprocessed predictions) and ``targets``.
        """
        samples, targets = batch
        outputs = self._resolve_eval_model()(samples)
        if self.train_config.compute_val_loss:
            loss_dict = self.criterion(outputs, targets)
            weight_dict = self.criterion.weight_dict
            loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict if k in weight_dict)
            self._log_val_loss_metrics(loss, loss_dict, batch_size=len(targets))

        orig_sizes = torch.stack([t["orig_size"] for t in targets])
        results = self.postprocess(outputs, orig_sizes)
        return {"results": results, "targets": targets}

    @property
    def _fused_adamw_env_eligible(self) -> bool:
        """Return whether the runtime would enable fused AdamW, ignoring optimizer choice.

        Captures only the hardware/precision preconditions (BF16 on CUDA), so the
        custom-optimizer path can tell whether a dropped ``fused_optimizer=True``
        would actually have mattered.

        Returns:
            ``True`` when fused AdamW is requested and the runtime supports it.
        """
        return (
            self.model_config.fused_optimizer
            and torch.cuda.is_available()
            and torch.cuda.is_bf16_supported()
            and str(self.trainer.precision) in {"bf16-mixed", "bf16", "bf16-true"}
        )

    @property
    def _use_fused_optimizer(self) -> bool:
        """Return whether fused AdamW should be used for the current training configuration.

        Fused AdamW is only safe when the trainer's actual precision is a BF16 variant.  Checking GPU capability alone
        (``is_bf16_supported()``) is
        insufficient: on Ampere+ hardware that flag is always ``True`` even when
        the trainer is configured for ``32-true``, which causes a ``params, grads, exp_avgs, and exp_avg_sqs must have
        same dtype, device, and layout`` crash in DDP because gradient bucket views have non-matching strides in FP32.
        It additionally requires the built-in ``optimizer="adamw"`` selection: fused state normalization and gradient
        clipping are AdamW-specific and must not fire for custom optimizers.

        Returns:
            ``True`` when fused AdamW is requested, safe, and the built-in AdamW optimizer is selected.

        Examples:
            >>> from unittest.mock import patch
            >>> module = RFDETRModelModule.__new__(RFDETRModelModule)
            >>> module.model_config = type("Cfg", (), {"fused_optimizer": True})()
            >>> with patch("torch.cuda.is_available", return_value=False):
            ...     module._use_fused_optimizer
            False
        """
        if not self._fused_adamw_env_eligible:
            return False
        return _is_builtin_fused_adamw(self.train_config.optimizer)

    def configure_optimizers(self) -> OptimizerLRSchedulerConfig:
        """Build the configured optimizer with layer-wise LR decay and scheduler.

        Uses ``trainer.estimated_stepping_batches`` for total step count so cosine annealing covers the full training
        run regardless of dataset size or accumulation settings.
        ``optimizer="adamw"`` keeps RF-DETR's fused torch AdamW path;
        other names can be loaded from ``pytorch-optimizer``.

        Returns:
            PTL optimizer config dict with optimizer and step-interval scheduler.
        """
        tc = self.train_config
        ns = _namespace_from_configs(self.model_config, tc)

        # Unwrap torch.compile's OptimizedModule so get_param_dict sees the
        # original module's named_parameters() — compiled wrapper can cause
        # name-prefix mismatches that put the same tensor in multiple groups.
        model_for_params = getattr(self.model, "_orig_mod", self.model)
        param_dicts = get_param_dict(ns, model_for_params)
        param_dicts = [param_group for param_group in param_dicts if param_group["params"].requires_grad]

        optimizer_cfg = tc.optimizer
        optimizer: torch.optim.Optimizer
        if _is_builtin_fused_adamw(optimizer_cfg):
            # Built-in managed fused AdamW path (unchanged behavior).
            try:
                optimizer = torch.optim.AdamW(
                    param_dicts,
                    lr=tc.lr,
                    weight_decay=tc.weight_decay,
                    fused=self._use_fused_optimizer,
                    **tc.optimizer_kwargs,
                )
            except TypeError as exc:
                raise TypeError(
                    f"Failed to initialize optimizer 'adamw': {exc}. "
                    "Check optimizer_kwargs for arguments supported by torch.optim.AdamW."
                ) from exc
        else:
            if self._fused_adamw_env_eligible:
                logger.warning(_FUSED_IGNORED_MSG, optimizer_cfg)
            if not isinstance(optimizer_cfg, str):
                # Explicit callable / functools.partial: called with param groups only.
                callable_name = getattr(optimizer_cfg, "__qualname__", None) or repr(optimizer_cfg)
                optimizer = _instantiate_explicit_optimizer(optimizer_cfg, callable_name, param_dicts, {})
            elif _is_managed_optimizer_name(optimizer_cfg):
                # Managed native torch.optim short name (lr + signature-aware weight_decay injected).
                native_class: _OptimizerFactory = _resolve_native_optimizer(optimizer_cfg)
                optimizer = _instantiate_optimizer(native_class, optimizer_cfg, param_dicts, tc)
            else:
                # Explicit dotted import path: constructed from optimizer_kwargs only.
                optimizer_class = _import_optimizer_class(optimizer_cfg)
                optimizer = _instantiate_explicit_optimizer(
                    optimizer_class, optimizer_cfg, param_dicts, tc.optimizer_kwargs
                )

        # ``trainer.estimated_stepping_batches`` is reported in *microbatch* units when
        # the keypoint path runs with ``Trainer(accumulate_grad_batches=1)`` and manages
        # accumulation manually. ``LambdaLR.step()`` is called once per optimizer-step
        # (i.e. every ``grad_accum_steps`` microbatches), so the schedule must be sized
        # in optimizer-step units rather than microbatches; otherwise warmup and cosine
        # decay finish ``grad_accum_steps``× too early. Detection / segmentation models
        # still rely on Lightning's automatic optimization, where PTL already accounts
        # for ``accumulate_grad_batches`` inside ``estimated_stepping_batches`` and the
        # division below is a no-op (``grad_accum_steps`` would be 1 in that path).
        grad_accum_steps = max(1, int(tc.grad_accum_steps))
        microbatches = int(self.trainer.estimated_stepping_batches)
        # _should_step_optimizer steps the final partial window at epoch end, so the true
        # number of optimizer steps is ceil(microbatches / grad_accum_steps).  Using floor
        # would undercount when the epoch is not evenly divisible, causing warmup / cosine
        # schedules to finish one step earlier than the last actual step fires.
        total_steps = (
            max(1, math.ceil(microbatches / grad_accum_steps)) if self._use_manual_optimization else microbatches
        )
        steps_per_epoch = max(1, total_steps // tc.epochs)
        warmup_steps = int(steps_per_epoch * tc.warmup_epochs)

        scheduler_cfg = tc.lr_scheduler
        scheduler: LRScheduler | ReduceLROnPlateau
        interval = "step"
        monitor: str | None = None
        if _is_managed_scheduler_name(scheduler_cfg):
            # Managed "step" / "cosine" preset — warmup + total-step sizing baked into a LambdaLR (unchanged behavior).
            scheduler = _build_managed_scheduler(optimizer, tc, total_steps, steps_per_epoch, warmup_steps)
        else:
            if not isinstance(scheduler_cfg, str):
                # Explicit callable / functools.partial: built from the optimizer only (kwargs baked in).
                scheduler_name = getattr(scheduler_cfg, "__qualname__", None) or repr(scheduler_cfg)
                scheduler = _instantiate_explicit_scheduler(scheduler_cfg, scheduler_name, optimizer, {})
            else:
                # Explicit dotted import path: constructed from lr_scheduler_kwargs only.
                scheduler_class = _import_scheduler_class(scheduler_cfg)
                scheduler = _instantiate_explicit_scheduler(
                    scheduler_class, scheduler_cfg, optimizer, tc.lr_scheduler_kwargs
                )
            interval = tc.lr_scheduler_interval
            if isinstance(scheduler, ReduceLROnPlateau):
                monitor = tc.lr_scheduler_monitor
                # The monitored metric (e.g. val/loss) is only available per epoch, so plateau always steps
                # on the epoch boundary regardless of the configured interval.
                interval = "epoch"
                if warmup_steps > 0:
                    logger.warning(
                        "warmup_epochs=%s is ignored for ReduceLROnPlateau; a metric-driven scheduler cannot "
                        "be composed with a linear warmup ramp.",
                        tc.warmup_epochs,
                    )
            else:
                # Auto-wrap explicit schedulers with a linear warmup ramp. The wrap is stepped at the same cadence as
                # the scheduler, so size it in the scheduler's own units: optimizer steps for "step", epochs for "epoch"
                # (otherwise a step-sized ramp stepped once per epoch would stretch across the whole run).
                if interval == "step":
                    warmup_units = warmup_steps
                else:
                    # Epoch cadence: the ramp is stepped once per epoch. ceil keeps a fractional warmup_epochs from
                    # truncating to zero (a silently dropped warmup). A single-epoch ramp has start_factor == 1.0,
                    # i.e. a flat no-op that looks like warmup but isn't, so a gradual epoch-granular warmup needs
                    # >= 2 epochs; warn and skip rather than emit a degenerate ramp.
                    warmup_units = math.ceil(tc.warmup_epochs)
                    if tc.warmup_epochs > 0 and warmup_units < 2:
                        logger.warning(
                            "warmup_epochs=%s with lr_scheduler_interval='epoch' cannot form a gradual warmup ramp "
                            "(epoch-granular warmup needs >= 2 epochs); skipping warmup. Use "
                            "lr_scheduler_interval='step' for sub-epoch warmup, or set warmup_epochs >= 2.",
                            tc.warmup_epochs,
                        )
                        warmup_units = 0
                if warmup_units > 0:
                    scheduler = _wrap_with_warmup(scheduler, optimizer, warmup_units)

        self._lr_scheduler_interval = interval
        self._lr_scheduler_monitor = monitor

        lr_scheduler_config: LRSchedulerConfigType = {"scheduler": scheduler, "interval": interval}
        if monitor is not None:
            lr_scheduler_config["monitor"] = monitor
        return {
            "optimizer": optimizer,
            "lr_scheduler": lr_scheduler_config,
        }

    def clip_gradients(
        self,
        optimizer: torch.optim.Optimizer | LightningOptimizer,
        gradient_clip_val: float | None = None,
        gradient_clip_algorithm: str | None = None,
    ) -> None:
        """Override PTL gradient clipping to support fused AdamW.

        PTL's AMP precision plugin refuses to clip gradients when the optimizer declares it handles unscaling internally
        (fused=True).  When fused is active we are on BF16 (no GradScaler) so ``clip_grad_norm_`` is correct.  For the
        non-fused path (FP16 + GradScaler or FP32) we delegate to ``super()`` to preserve scaler-aware unscaling.

        Args:
            optimizer: The current optimizer.
            gradient_clip_val: Maximum gradient norm.
            gradient_clip_algorithm: Clipping algorithm; forwarded to super()
                for the non-fused path.
        """
        if self._use_fused_optimizer:
            if gradient_clip_val and gradient_clip_val > 0:
                torch.nn.utils.clip_grad_norm_(self.parameters(), gradient_clip_val)
        else:
            # PTL's own type stub only declares Optimizer here, but LightningOptimizer dynamically
            # multiply-inherits from the wrapped optimizer's class, so it satisfies this at runtime too.
            super().clip_gradients(
                optimizer,  # type: ignore[arg-type]
                gradient_clip_val=gradient_clip_val,
                gradient_clip_algorithm=gradient_clip_algorithm,
            )

    @staticmethod
    def _normalize_optimizer_state(optimizer: torch.optim.Optimizer) -> int:
        """Cast restored floating-point optimizer state tensors to match live parameters.

        Args:
            optimizer: AdamW optimizer whose state may have been rehydrated with a mismatched dtype or layout.

        Returns:
            Number of state tensors that were reallocated to match the current parameter layout.
        """
        normalized = 0
        for group in optimizer.param_groups:
            for param in group["params"]:
                state = optimizer.state.get(param)
                if not state:
                    continue
                for key, value in list(state.items()):
                    if not isinstance(value, torch.Tensor) or not value.is_floating_point():
                        continue
                    if value.shape != param.shape:
                        continue
                    if value.device == param.device and value.dtype == param.dtype and value.stride() == param.stride():
                        continue
                    restored = torch.empty_like(param)
                    restored.copy_(value.to(device=param.device, dtype=param.dtype))
                    state[key] = restored
                    normalized += 1
        return normalized

    def test_step(self, batch: tuple[Any, Any], batch_idx: int) -> dict[str, Any]:
        """Run forward pass and postprocess for one test step.

        Mirrors :meth:`validation_step` so ``COCOEvalCallback`` can accumulate results via ``on_test_batch_end`` when
        ``trainer.test()`` is called (e.g. from :class:`~rfdetr.training.callbacks.BestModelCallback` at end of
        training).

        Args:
            batch: Tuple of (NestedTensor samples, list of target dicts).
            batch_idx: Batch index within the test epoch.

        Returns:
            Dict with ``results`` (postprocessed predictions) and ``targets``.
        """
        samples, targets = batch
        outputs = self.model(samples)
        if self.train_config.compute_test_loss:
            loss_dict = self.criterion(outputs, targets)
            weight_dict = self.criterion.weight_dict
            loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict if k in weight_dict)
            self.log("test/loss", loss, sync_dist=True, batch_size=len(targets))

        orig_sizes = torch.stack([t["orig_size"] for t in targets])
        results = self.postprocess(outputs, orig_sizes)
        return {"results": results, "targets": targets}

    def predict_step(self, batch: tuple[Any, Any], batch_idx: int, dataloader_idx: int = 0) -> Any:
        """Run inference on a preprocessed batch and return postprocessed results.

        Args:
            batch: Tuple of (NestedTensor samples, list of target dicts).
            batch_idx: Batch index.
            dataloader_idx: Index of the predict dataloader.

        Returns:
            Postprocessed detection results from ``PostProcess``.
        """
        samples, targets = batch
        with torch.no_grad():
            outputs = self.model(samples)
        orig_sizes = torch.stack([t["orig_size"] for t in targets])
        return self.postprocess(outputs, orig_sizes)

    def on_load_checkpoint(self, checkpoint: dict[str, Any]) -> None:
        """Auto-detect legacy formats and reconcile PE shapes at checkpoint load time.

        PTL calls this hook before applying ``checkpoint["state_dict"]`` to the module.  Three normalisation steps are
        applied in order:

        1. **Raw legacy format** — a ``*.pth`` file loaded directly by
           ``Trainer`` (e.g. via ``ckpt_path=``).  Recognised by the presence of ``"model"`` without ``"state_dict"``.
           The state dict is rewritten in-place with the ``"model."`` prefix so PTL can apply it normally.

        2. **Positional-embedding interpolation** — when the checkpoint was
           saved at a different image resolution than the current model, the DINOv2 ``position_embeddings`` tensor shape
           will mismatch. :func:`~rfdetr.models.weights.interpolate_position_embeddings` is called to bicubic-resize the
           PE to ``model_config.positional_encoding_size`` before PTL applies the state dict.  Regression fix for
           :issue:`998`.

        3. **Converted format** — a file produced by
           :func:`~rfdetr.training.checkpoint.convert_legacy_checkpoint` that already has ``"state_dict"`` but also
           carries ``"legacy_ema_state_dict"``.  The EMA weights are stashed on ``self._pending_legacy_ema_state`` for
           optional restoration by :class:`~rfdetr.training.callbacks.ema.RFDETREMACallback`.

        Note:
            This hook only fires on ``Trainer(ckpt_path=...)`` resume paths. Fresh-train bootstrap from a
            ``pretrain_weights`` checkpoint runs through :func:`~rfdetr.models.weights.load_pretrain_weights` during
            ``__init__`` instead — that helper performs its own PTL ``.ckpt`` normalisation (``state_dict`` → ``model``
            key, ``_orig_mod`` strip) and PE interpolation, so the two code paths intentionally do not share state.

        Args:
            checkpoint: Checkpoint dict passed in by PTL (mutated in-place).
        """
        # Raw legacy .pth: no "state_dict" key — build it from "model".
        if "model" in checkpoint and "state_dict" not in checkpoint:
            checkpoint["state_dict"] = {"model." + k: v for k, v in checkpoint["model"].items()}

        # Interpolate DINOv2 positional embeddings when the checkpoint was saved
        # at a different resolution than the current model.  PTL applies
        # checkpoint["state_dict"] immediately after this hook, so the shapes
        # must already match at this point.  Regression: #998.
        if "state_dict" in checkpoint:
            interpolate_position_embeddings(
                checkpoint["state_dict"],
                self.model_config.positional_encoding_size,
            )

        # Stash legacy EMA weights for RFDETREMACallback.setup(), which restores
        # them into AveragedModel when resuming from converted legacy checkpoints.
        if "legacy_ema_state_dict" in checkpoint:
            self._pending_legacy_ema_state = checkpoint["legacy_ema_state_dict"]
            warnings.warn(
                "Checkpoint contains legacy EMA weights (`legacy_ema_state_dict`). "
                "Add RFDETREMACallback to your trainer callbacks to restore them; "
                "without it the stashed weights will be ignored.",
                UserWarning,
                stacklevel=2,
            )

    def reinitialize_detection_head(self, num_classes: int) -> None:
        """Reinitialize the detection head for a new class count.

        Args:
            num_classes: New number of classes (excluding background).
        """
        self.model.reinitialize_detection_head(num_classes)
