"""Utilities for optimizing inference performance in TrainTrack."""

import torch
from typing import Optional
import logging

logger = logging.getLogger(__name__)


def detect_optimal_batch_size(
    model: torch.nn.Module,
    max_batch_size: int = 16,
    min_batch_size: int = 1,
) -> int:
    """
    Auto-detect optimal batch size based on available GPU memory.
    
    Uses a simple heuristic: estimates model size and available memory,
    then suggests a batch size that should fit comfortably.
    
    Args:
        model: The model to analyze
        max_batch_size: Maximum batch size to consider
        min_batch_size: Minimum batch size to return
    
    Returns:
        Recommended batch size
    """
    try:
        device = next(model.parameters()).device
        
        # Only auto-detect for CUDA
        if device.type != 'cuda':
            logger.info("CPU detected - using batch_size=1")
            return 1
        
        # Get GPU memory info
        total_memory = torch.cuda.get_device_properties(device).total_memory
        reserved_memory = torch.cuda.memory_reserved(device)
        allocated_memory = torch.cuda.memory_allocated(device)
        free_memory = total_memory - max(reserved_memory, allocated_memory)
        
        # Estimate model size (rough approximation)
        param_count = sum(p.numel() for p in model.parameters())
        bytes_per_param = 2  # Assume fp16/bf16
        model_size_bytes = param_count * bytes_per_param
        
        # Conservative estimate: use 40% of free memory for batch
        # (rest needed for activations, gradients during training)
        usable_memory = free_memory * 0.4
        
        # Estimate memory per sample (very rough: 10x model size for a batch of 1)
        memory_per_sample = model_size_bytes * 10
        
        # Calculate batch size
        estimated_batch_size = int(usable_memory / memory_per_sample)
        batch_size = max(min_batch_size, min(estimated_batch_size, max_batch_size))
        
        logger.info(
            f"Auto-detected batch_size={batch_size} "
            f"(GPU: {total_memory / 1e9:.1f}GB total, "
            f"{free_memory / 1e9:.1f}GB free, "
            f"model: {param_count / 1e6:.1f}M params)"
        )
        
        return batch_size
        
    except Exception as e:
        logger.warning(f"Failed to auto-detect batch size: {e}. Using batch_size=4")
        return 4


def try_compile_model(model: torch.nn.Module, enable: bool = True) -> torch.nn.Module:
    """
    Attempt to compile model with torch.compile for faster inference.
    Falls back to uncompiled model if compilation fails or is disabled.
    
    Args:
        model: Model to compile
        enable: Whether to attempt compilation
    
    Returns:
        Compiled model (if successful) or original model
    """
    if not enable:
        return model
    
    try:
        import torch
        if hasattr(torch, 'compile'):
            logger.info("Compiling model with torch.compile() for faster inference...")
            compiled = torch.compile(model, mode='reduce-overhead')
            logger.info("✓ Model compiled successfully")
            return compiled
        else:
            logger.debug("torch.compile not available (PyTorch < 2.0)")
            return model
    except Exception as e:
        logger.warning(f"Failed to compile model: {e}. Using uncompiled model.")
        return model


def get_optimized_generation_config() -> dict:
    """
    Returns optimized generation kwargs for faster inference.
    
    These settings prioritize speed while maintaining quality:
    - use_cache=True: Faster autoregressive generation
    - do_sample=False: Greedy decoding (fastest)
    - num_beams=1: No beam search (faster)
    """
    return {
        "use_cache": True,
        "do_sample": False,
        "num_beams": 1,
    }
