from __future__ import annotations

import gc
import logging
import math
from typing import Any, Callable

import gin
import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


@gin.configurable
class HardwareConfig:
    """GPU hardware profile for the augmentation pipeline."""

    def __init__(
        self,
        profile_name: str = "rtx4090",
        total_ram_gb: float = 24.0,
        reserve_gb: float = 2.0,
        use_fp16: bool = True,
        batch_size: int | None = None,
        num_workers: int = 4,
    ) -> None:
        self.profile_name = profile_name
        self.total_ram_gb = total_ram_gb
        self.reserve_gb = reserve_gb
        self.use_fp16 = use_fp16
        self.batch_size = batch_size
        self.num_workers = num_workers
        # Derived.
        self.available_gb = self.total_ram_gb - self.reserve_gb

    def auto_batch_size(self, act_per_sample_mb: float,
                        overhead_mb: float = 200.0) -> int:
        """Compute safe batch size from actual free VRAM after model loading.

        If CUDA is available, reads real free VRAM. Otherwise falls back
        to config-based estimate.

        Returns manual batch_size if set in config.
        """
        if self.batch_size:
            return self.batch_size

        if torch.cuda.is_available():
            gpu_id = torch.cuda.current_device()
            free_driver, _ = torch.cuda.mem_get_info(gpu_id)
            cache_free = (torch.cuda.memory_reserved(gpu_id)
                          - torch.cuda.memory_allocated(gpu_id))
            free_mb = (free_driver + cache_free) / (1024 * 1024) - overhead_mb
        else:
            free_mb = self.available_gb * 1024 - overhead_mb

        if free_mb <= 0:
            return 1
        max_b = int(free_mb / act_per_sample_mb)
        safe = max(1, 2 ** int(math.log2(max(1, int(max_b * 0.7)))))
        logger.info("🎮 auto_batch_size: %.0f MB free, %.0f MB/sample → batch=%d",
                     free_mb + overhead_mb, act_per_sample_mb, safe)
        return safe


    def find_batch_size(
        self,
        inference_fn: Callable[[torch.Tensor], Any],
        input_shape: tuple[int, ...],
        device: torch.device,
        dtype: torch.dtype = torch.float32,
        max_batch: int = 64,
        safety_factor: float = 0.70,
    ) -> int:
        """Find optimal batch size by measuring incremental VRAM per sample.

        Runs forward passes with batch=1 and batch=2 to isolate the true
        per-sample cost from one-time overhead (CUDA warmup, kernel JIT, etc.).
        Then calculates max batch from free VRAM with a safety margin.

        Args:
            inference_fn: Callable that accepts [B, C, H, W] tensor (on CPU)
                          and runs one forward pass.
            input_shape: (C, H, W) — shape of a single sample.
            device: CUDA device.
            dtype: Tensor dtype for the dummy input.
            max_batch: Upper bound to try.
            safety_factor: Fraction of free VRAM to actually use (0.0–1.0).

        Returns:
            Optimal batch size (power of 2, >= 1).
        """
        if self.batch_size:
            logger.info("  batch_size forced by config: %d", self.batch_size)
            return self.batch_size

        if device.type != "cuda" or not torch.cuda.is_available():
            logger.info("  No CUDA device — batch_size=1")
            return 1

        gpu_id = torch.cuda.current_device()

        # --- Step 1: warmup with batch=1 (absorbs one-time overhead) ---
        torch.cuda.empty_cache()
        gc.collect()
        dummy = torch.rand(1, *input_shape, dtype=dtype)
        inference_fn(dummy)
        torch.cuda.synchronize()
        del dummy
        torch.cuda.empty_cache()
        gc.collect()

        # --- Step 2: measure peak for batch=1 ---
        torch.cuda.reset_peak_memory_stats(gpu_id)
        dummy1 = torch.rand(1, *input_shape, dtype=dtype)
        inference_fn(dummy1)
        torch.cuda.synchronize()
        peak_bs1 = torch.cuda.max_memory_allocated(gpu_id)
        del dummy1
        torch.cuda.empty_cache()
        gc.collect()

        # --- Step 3: measure peak for batch=2 ---
        torch.cuda.reset_peak_memory_stats(gpu_id)
        dummy2 = torch.rand(2, *input_shape, dtype=dtype)
        inference_fn(dummy2)
        torch.cuda.synchronize()
        peak_bs2 = torch.cuda.max_memory_allocated(gpu_id)
        del dummy2
        torch.cuda.empty_cache()
        gc.collect()

        # --- Step 4: incremental per-sample cost ---
        per_sample_bytes = peak_bs2 - peak_bs1

        # If peak doesn't grow with batch (model processes images
        # sequentially and manages its own memory), VRAM is constant
        # regardless of batch size → use max_batch for prefetch benefit.
        tolerance = 10 * 1024 * 1024  # 10 MB noise margin
        if per_sample_bytes <= tolerance:
            logger.info(
                "🎮 find_batch_size: peak(bs=1)=%.0f MB, peak(bs=2)=%.0f MB — "
                "VRAM constant, using max_batch=%d",
                peak_bs1 / (1024 * 1024), peak_bs2 / (1024 * 1024), max_batch,
            )
            return max_batch

        # --- Step 5: compute batch size from free VRAM ---
        free_driver, _ = torch.cuda.mem_get_info(gpu_id)
        cache_free = (torch.cuda.memory_reserved(gpu_id)
                      - torch.cuda.memory_allocated(gpu_id))
        usable_bytes = (free_driver + cache_free) * safety_factor

        per_sample_mb = per_sample_bytes / (1024 * 1024)
        usable_mb = usable_bytes / (1024 * 1024)

        logger.info(
            "🎮 find_batch_size: peak(bs=1)=%.0f MB, peak(bs=2)=%.0f MB, "
            "incremental=%.1f MB/sample",
            peak_bs1 / (1024 * 1024), peak_bs2 / (1024 * 1024),
            per_sample_mb,
        )

        if per_sample_bytes <= 0:
            logger.warning("  per-sample VRAM ≤ 0 (%.1f MB), defaulting to batch=1",
                           per_sample_mb)
            return 1

        max_b = int(usable_bytes / per_sample_bytes)
        # Round down to power of 2 for stable GPU occupancy.
        best = max(1, min(max_batch, 2 ** int(math.log2(max(1, max_b)))))

        logger.info(
            "🎮 find_batch_size: %.0f MB free × %.0f%% = %.0f MB usable, "
            "%.1f MB/sample → batch=%d",
            (free_driver + cache_free) / (1024 * 1024),
            safety_factor * 100,
            usable_mb,
            per_sample_mb,
            best,
        )
        return best


def get_hardware_cfg(path2cfg: str) -> HardwareConfig:
    """Load ONLY hardware config (for isolated testing)."""
    gin.clear_config()
    gin.parse_config_file(f"{path2cfg}hardware.gin")
    return HardwareConfig()
