"""VRAM usage estimation for training configurations."""

from __future__ import annotations

from dataclasses import dataclass

from bilgeai.optimizer.models import TrainingConfig

# Approximate model sizes in billions of parameters
MODEL_SIZES: dict[str, float] = {
    "llama-3-8b": 8.0,
    "llama-3-70b": 70.0,
    "mistral-7b": 7.0,
    "qwen-7b": 7.0,
    "qwen-14b": 14.0,
    "qwen-72b": 72.0,
}


@dataclass
class VRAMEstimate:
    """Estimated VRAM usage breakdown."""

    model_weights_gb: float
    activations_gb: float
    optimizer_states_gb: float
    total_gb: float
    fits_in_vram: bool
    available_vram_gb: float


def estimate_vram(config: TrainingConfig, vram_gb: float) -> VRAMEstimate:
    """Estimate VRAM usage for a training configuration."""
    # Guess model size from name
    model_params_b = _guess_model_params(config.model_name)

    # Bytes per parameter based on quantization
    quant_bytes = {"4bit": 0.5, "8bit": 1.0, "none": 2.0}.get(config.quantization or "4bit", 0.5)

    # Model weights
    model_weights_gb = model_params_b * quant_bytes / 1024

    # Activation memory (rough estimate based on batch_size and seq_length)
    activation_per_sample = (config.max_seq_length / 512) * 0.3  # GB per sample
    activations_gb = activation_per_sample * config.batch_size

    # Gradient checkpointing reduces activation memory by ~60%
    if config.use_gradient_checkpointing:
        activations_gb *= 0.4

    # LoRA optimizer states are small (only adapter weights)
    if config.lora_enabled and config.lora_rank:
        lora_params = model_params_b * 0.01 * (config.lora_rank / 16)  # rough
        optimizer_states_gb = lora_params * 8 / 1024  # AdamW: 8 bytes per param
    else:
        optimizer_states_gb = model_params_b * 8 / 1024

    total_gb = model_weights_gb + activations_gb + optimizer_states_gb

    return VRAMEstimate(
        model_weights_gb=round(model_weights_gb, 2),
        activations_gb=round(activations_gb, 2),
        optimizer_states_gb=round(optimizer_states_gb, 2),
        total_gb=round(total_gb, 2),
        fits_in_vram=total_gb <= vram_gb * 0.95,  # 5% safety margin
        available_vram_gb=vram_gb,
    )


def _guess_model_params(model_name: str) -> float:
    """Guess model size in billions from the name."""
    name_lower = model_name.lower()
    for key, size in MODEL_SIZES.items():
        if key in name_lower:
            return size

    # Try to extract number from name (e.g., "7b", "13b", "70b")
    import re

    match = re.search(r"(\d+)[bB]", model_name)
    if match:
        return float(match.group(1))

    return 7.0  # default assumption
