"""Auto-batch sizing for GPU inference.

ControlMT's `translate()` runs beam search sentence-by-sentence, but we
ThreadPoolExecutor across N concurrent calls. N is bounded by free VRAM.

Memory per concurrent translation at beam=2, max_len=256:
    fp16: ~25 MB    (encoder + decoder state + beam tensors + kv-cache)
    bf16: ~25 MB
    fp32: ~50 MB
"""

from __future__ import annotations
import warnings


_MEM_PER_SENT_MB = {
    "float16":  25,
    "bfloat16": 25,
    "float32":  50,
    "int8-dynamic": 12,
}
_MAX_BATCH = 64       # ControlMT.translate() beam-search per-sentence — gain plateaus past ~16
_TARGET_VRAM_FRAC = 0.8


def auto_batch_size(device: str, dtype: str, quant: str = "none",
                    free_vram_mb: float | None = None) -> int:
    """Estimate batch size that fits in ~80% of free VRAM.

    Returns 1 on CPU (no auto-batching there) with a warning.
    On GPU returns 1..MAX_BATCH.
    """
    if device != "cuda":
        warnings.warn(
            "auto_batch=True ignored on CPU (no VRAM signal to size against). "
            "Pass an explicit batch_size for batched CPU translation.",
            RuntimeWarning, stacklevel=2)
        return 1

    if free_vram_mb is None:
        import torch
        free, _total = torch.cuda.mem_get_info()
        free_vram_mb = free / 1024**2

    per_sent_mb = _MEM_PER_SENT_MB.get(quant if quant != "none" else dtype, 50)
    n = int((free_vram_mb * _TARGET_VRAM_FRAC) // per_sent_mb)
    return max(1, min(_MAX_BATCH, n))