# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
"""Auto-batch probing: find a safe micro-batch size before training.

Probe assumptions (worst-case so training does not OOM):
- Resolution: When multi_scale is True we use the maximum of the multi-scale
  augmentation scales (same as compute_multi_scale_scales). Otherwise we use model resolution. This ensures the step
  uses the max resolution seen in training.
- Targets: Memory grows with number of targets per image. We use
  auto_batch_max_targets_per_image (config) to synthesize that many targets per image so the probe reflects worst-case
  matcher and loss memory.
- EMA: When use_ema is True, an EMA copy of the model is kept in memory. We
  apply auto_batch_ema_headroom (e.g. 0.7) to the probed batch size so the effective safe batch leaves room for the EMA
  model.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Any, cast

import torch

from rfdetr.config import ModelConfig, TrainConfig
from rfdetr.datasets.coco import compute_multi_scale_scales
from rfdetr.models import build_criterion_from_config
from rfdetr.utilities.logger import get_logger
from rfdetr.utilities.tensors import NestedTensor

logger = get_logger()


@dataclass(frozen=True)
class AutoBatchResult:
    """Result of auto-batch probing: safe micro-batch size and recommended grad accumulation.

    Attributes:
        safe_micro_batch: Per-device batch size that fits in memory for one train step.
        recommended_grad_accum_steps: Steps to accumulate to reach target effective batch.
        effective_batch_size: safe_micro_batch * recommended_grad_accum_steps.
        device_name: Human-readable GPU name used for probing.
    """

    safe_micro_batch: int
    recommended_grad_accum_steps: int
    effective_batch_size: int
    device_name: str


def _is_cuda_oom(exc: BaseException) -> bool:
    message = str(exc).lower()
    return "out of memory" in message or "cuda error: out of memory" in message


def _make_synthetic_batch(
    micro_batch_size: int,
    resolution: int,
    device: torch.device,
    num_classes: int,
    segmentation_head: bool = False,
    max_targets_per_image: int = 1,
    num_channels: int = 3,
) -> tuple[NestedTensor, list[dict[str, torch.Tensor]]]:
    """Build a minimal (samples, targets) batch for probing.

    Uses max_targets_per_image targets per image so memory reflects worst-case matcher and loss. When segmentation_head
    is True, each target dict includes "masks" of shape (max_targets_per_image, resolution, resolution).
    """
    tensors = torch.randn(micro_batch_size, num_channels, resolution, resolution, device=device)
    mask = torch.zeros(micro_batch_size, resolution, resolution, dtype=torch.bool, device=device)
    samples = NestedTensor(tensors, mask)

    max_label = max(0, num_classes - 1)
    n = max(1, max_targets_per_image)
    targets: list[dict[str, torch.Tensor]] = []
    for idx in range(micro_batch_size):
        # Replicate one box/label n times so matcher and loss see n targets per image.
        boxes = torch.tensor([[0.5, 0.5, 0.2, 0.2]], dtype=torch.float32, device=device).expand(n, 4)
        labels = torch.tensor([min(1, max_label)], dtype=torch.int64, device=device).expand(n)
        iscrowd = torch.zeros(n, dtype=torch.int64, device=device)
        area = torch.full((n,), 0.04, dtype=torch.float32, device=device)
        t: dict[str, torch.Tensor] = {
            "boxes": boxes,
            "labels": labels,
            "image_id": torch.tensor(idx, dtype=torch.int64, device=device),
            "orig_size": torch.tensor([resolution, resolution], dtype=torch.int64, device=device),
            "size": torch.tensor([resolution, resolution], dtype=torch.int64, device=device),
            "iscrowd": iscrowd,
            "area": area,
        }
        if segmentation_head:
            t["masks"] = torch.zeros(n, resolution, resolution, dtype=torch.bool, device=device)
        targets.append(t)
    return samples, targets


def _probe_step(
    model: torch.nn.Module,
    criterion: torch.nn.Module,
    micro_batch_size: int,
    resolution: int,
    device: torch.device,
    num_classes: int,
    amp: bool,
    segmentation_head: bool = False,
    max_targets_per_image: int = 1,
    num_channels: int = 3,
    autocast_dtype: torch.dtype | None = None,
) -> bool:
    """Run one forward + loss + backward; return True if successful, False on OOM."""
    try:
        model.zero_grad(set_to_none=True)
        criterion.zero_grad(set_to_none=True)
        samples, targets = _make_synthetic_batch(
            micro_batch_size=micro_batch_size,
            resolution=resolution,
            device=device,
            num_classes=num_classes,
            segmentation_head=segmentation_head,
            max_targets_per_image=max_targets_per_image,
            num_channels=num_channels,
        )

        autocast_kwargs: dict[str, Any] = {"device_type": "cuda", "enabled": amp}
        if autocast_dtype is not None:
            autocast_kwargs["dtype"] = autocast_dtype
        with torch.autocast(**autocast_kwargs):
            outputs = model(samples, targets)
            loss_dict = cast("dict[str, torch.Tensor]", criterion(outputs, targets))
            weight_dict = cast("dict[str, float]", criterion.weight_dict)
            weighted_losses = [loss_dict[name] * weight_dict[name] for name in loss_dict if name in weight_dict]
            if not weighted_losses:
                raise RuntimeError(
                    "auto-batch probe could not build weighted losses: no overlap between criterion loss_dict and "
                    "weight_dict keys.",
                )
            loss = torch.stack(weighted_losses).sum()

        if not torch.isfinite(loss):
            raise RuntimeError("auto-batch probe produced a non-finite training loss.")

        torch.autograd.backward(loss)
        model.zero_grad(set_to_none=True)
        criterion.zero_grad(set_to_none=True)
        return True
    except RuntimeError as exc:
        if _is_cuda_oom(exc):
            torch.cuda.empty_cache()
            return False
        raise


def probe_max_micro_batch(
    model: torch.nn.Module,
    criterion: torch.nn.Module,
    resolution: int,
    device: torch.device,
    num_classes: int,
    amp: bool,
    segmentation_head: bool = False,
    max_targets_per_image: int = 1,
    safety_margin: float = 0.9,
    max_micro_batch: int = 128,
    num_channels: int = 3,
    autocast_dtype: torch.dtype | None = None,
) -> int:
    """Find the largest per-device batch size that fits in memory for one train step.

    Uses exponential search (1, 2, 4, ...) up to the first failure, then binary search between the last successful size
    and the first failure to get the exact maximum. The returned value is floor(max_ok * safety_margin), so
    safety_margin in (0, 1] scales down the result for headroom (e.g. 0.9 keeps 10% margin).

    Args:
        model: The model to probe (will be set to train mode).
        criterion: The loss criterion (must match model output and target format).
        resolution: Input spatial size (square).
        device: CUDA device to run on.
        num_classes: Number of classes (for synthetic targets).
        amp: Whether to use autocast for the forward.
        segmentation_head: If True, synthetic targets include "masks" for loss_masks.
        max_targets_per_image: Number of synthetic targets per image (worst-case memory).
        safety_margin: Fraction of max batch to return (0 < safety_margin <= 1).
        max_micro_batch: Cap on batch size to try.
        num_channels: Number of input image channels (for synthetic probe images).

    Returns:
        Safe micro-batch size (>= 1).

    Raises:
        RuntimeError: If device is not CUDA or if micro_batch_size=1 already fails (OOM).
        ValueError: If max_micro_batch < 1 or safety_margin not in (0, 1].
    """
    if device.type != "cuda":
        raise RuntimeError("auto-batch probing currently supports CUDA only.")
    if max_micro_batch < 1:
        raise ValueError("max_micro_batch must be >= 1.")
    if not (0 < safety_margin <= 1.0):
        raise ValueError("safety_margin must be in (0, 1].")

    model_training = model.training
    criterion_training = criterion.training
    model.train()
    criterion.train()

    try:
        lower_ok = 0
        candidate = 1
        upper_fail = None

        while candidate <= max_micro_batch:
            if _probe_step(
                model,
                criterion,
                candidate,
                resolution,
                device,
                num_classes,
                amp,
                segmentation_head,
                max_targets_per_image,
                num_channels,
                autocast_dtype,
            ):
                lower_ok = candidate
                candidate *= 2
            else:
                upper_fail = candidate
                break

        if lower_ok < 1:
            raise RuntimeError(
                "auto-batch probe failed at micro_batch_size=1. "
                "Try lowering resolution or enabling gradient_checkpointing."
            )

        if upper_fail is None:
            upper_fail = max_micro_batch + 1

        lo = lower_ok + 1
        hi = min(upper_fail - 1, max_micro_batch)
        while lo <= hi:
            mid = (lo + hi) // 2
            if _probe_step(
                model,
                criterion,
                mid,
                resolution,
                device,
                num_classes,
                amp,
                segmentation_head,
                max_targets_per_image,
                num_channels,
                autocast_dtype,
            ):
                lower_ok = mid
                lo = mid + 1
            else:
                hi = mid - 1

        # safe_micro_batch <= lower_ok always, since safety_margin <= 1.0.
        safe_micro_batch = max(1, math.floor(lower_ok * safety_margin))
        return safe_micro_batch
    finally:
        model.train(model_training)
        criterion.train(criterion_training)
        model.zero_grad(set_to_none=True)
        criterion.zero_grad(set_to_none=True)
        torch.cuda.empty_cache()


def recommend_grad_accum_steps(safe_micro_batch: int, target_effective_batch: int) -> int:
    """Recommend gradient accumulation steps to reach target effective batch size.

    Args:
        safe_micro_batch: Per-step batch size that fits in memory.
        target_effective_batch: Desired effective batch (micro_batch * accum_steps).

    Returns:
        ceil(target_effective_batch / safe_micro_batch), at least 1.

    Raises:
        ValueError: If either argument is < 1.
    """
    if safe_micro_batch < 1:
        raise ValueError("safe_micro_batch must be >= 1.")
    if target_effective_batch < 1:
        raise ValueError("target_effective_batch must be >= 1.")
    return max(1, math.ceil(target_effective_batch / safe_micro_batch))


def resolve_auto_batch_config(
    model_context: Any,
    model_config: ModelConfig,
    train_config: TrainConfig,
    safety_margin: float = 0.9,
    max_micro_batch: int = 128,
) -> AutoBatchResult:
    """Resolve batch_size='auto' into concrete batch_size and grad_accum_steps using a probe.

    Expects model_context to have attributes: .device (torch.device) and .model (nn.Module). Runs probe_max_micro_batch
    on the current model/criterion, then recommend_grad_accum_steps using train_config.auto_batch_target_effective. Logs
    device, segmentation flag, resolution, and the chosen values; also logs that the probe is train-step-only and that
    eval/test may use more memory.

    Args:
        model_context: Object with .device and .model (e.g. RFDETR.model from get_model()).
        model_config: Architecture config (resolution, num_classes, amp, segmentation_head).
        train_config: Training config (auto_batch_target_effective); batch_size should be "auto".
        safety_margin: Fraction of max batch to use (passed to probe_max_micro_batch).
        max_micro_batch: Upper bound on batch size to try (passed to probe_max_micro_batch).

    Returns:
        AutoBatchResult with safe_micro_batch, recommended_grad_accum_steps, effective_batch_size, and device_name.

    Raises:
        RuntimeError: If CUDA is not available or model_context.device is not CUDA.
    """
    device = model_context.device
    if not torch.cuda.is_available() or device.type != "cuda":
        raise RuntimeError("batch_size='auto' requires a CUDA device for probing in v1.")

    # Use max multi-scale resolution when multi_scale is True so probe reflects worst-case.
    multi_scale = getattr(train_config, "multi_scale", False)
    do_random_resize = getattr(train_config, "do_random_resize_via_padding", False)
    if multi_scale and not do_random_resize:
        expanded_scales = getattr(train_config, "expanded_scales", True)
        patch_size = getattr(model_config, "patch_size", 14)
        num_windows = getattr(model_config, "num_windows", 4)
        scales = compute_multi_scale_scales(
            model_config.resolution,
            expanded_scales,
            patch_size,
            num_windows,
        )
        probe_resolution = max(scales) if scales else model_config.resolution
    else:
        probe_resolution = model_config.resolution

    max_targets_per_image = getattr(train_config, "auto_batch_max_targets_per_image", 100)

    criterion, _ = build_criterion_from_config(model_config, train_config)
    criterion = criterion.to(device)

    amp_enabled = bool(model_config.amp)
    amp_dtype_str = getattr(train_config, "amp_dtype", "auto")
    if amp_enabled:
        if amp_dtype_str == "fp16":
            probe_autocast_dtype: torch.dtype | None = torch.float16
        else:
            # "bf16" or "auto" — both use bf16 on capable hardware, fp16 as fallback
            probe_autocast_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    else:
        probe_autocast_dtype = None

    safe_micro_batch = probe_max_micro_batch(
        model=model_context.model,
        criterion=criterion,
        resolution=probe_resolution,
        device=device,
        num_classes=model_config.num_classes,
        amp=amp_enabled,
        segmentation_head=model_config.segmentation_head,
        max_targets_per_image=max_targets_per_image,
        safety_margin=safety_margin,
        max_micro_batch=max_micro_batch,
        num_channels=getattr(model_config, "num_channels", 3),
        autocast_dtype=probe_autocast_dtype,
    )

    use_ema = getattr(train_config, "use_ema", False)
    if use_ema:
        headroom = getattr(train_config, "auto_batch_ema_headroom", 0.7)
        safe_micro_batch = max(1, math.floor(safe_micro_batch * headroom))
        logger.info("[auto-batch] Applied EMA headroom (%.2f): safe_micro_batch=%s", headroom, safe_micro_batch)

    # Infer world size from train configuration (only when explicit integers are provided)
    devices = getattr(train_config, "devices", None)
    num_nodes = getattr(train_config, "num_nodes", 1)
    if isinstance(devices, int) and isinstance(num_nodes, int):
        world_size = max(1, devices * num_nodes)
    else:
        world_size = 1

    # Interpret auto_batch_target_effective as a global target and derive a per-device target
    target_effective_global = train_config.auto_batch_target_effective
    if world_size > 1:
        target_effective_per_device = max(1, math.ceil(target_effective_global / world_size))
    else:
        target_effective_per_device = target_effective_global

    grad_accum_steps = recommend_grad_accum_steps(safe_micro_batch, target_effective_per_device)
    effective_batch_size_per_device = safe_micro_batch * grad_accum_steps
    global_effective_batch_size = effective_batch_size_per_device * world_size
    device_name = torch.cuda.get_device_name(device)

    logger.info(
        "[auto-batch] device=%s world_size=%s segmentation=%s probe_resolution=%s max_targets_per_image=%s",
        device_name,
        world_size,
        model_config.segmentation_head,
        probe_resolution,
        max_targets_per_image,
    )
    logger.info(
        "[auto-batch] safe_micro_batch=%s grad_accum_steps=%s effective_batch_per_device=%s global_effective_batch=%s",
        safe_micro_batch,
        grad_accum_steps,
        effective_batch_size_per_device,
        global_effective_batch_size,
    )
    logger.info("[auto-batch] This probe estimates train-step-safe micro-batch only.")
    logger.info("[auto-batch] Validation/test (especially segmentation mask eval) may require more memory.")

    return AutoBatchResult(
        safe_micro_batch=safe_micro_batch,
        recommended_grad_accum_steps=grad_accum_steps,
        effective_batch_size=effective_batch_size_per_device,
        device_name=device_name,
    )
