"""Tests for VRAM estimation."""

from bilgeai.optimizer.models import TrainingConfig
from bilgeai.validation.vram_estimator import _guess_model_params, estimate_vram


class TestGuessModelParams:
    def test_llama_8b(self):
        assert _guess_model_params("meta-llama/Llama-3-8B") == 8.0

    def test_mistral_7b(self):
        assert _guess_model_params("mistralai/Mistral-7B-v0.1") == 7.0

    def test_generic_number_extraction(self):
        assert _guess_model_params("some-model-13B") == 13.0

    def test_unknown_defaults_to_7(self):
        assert _guess_model_params("unknown-model") == 7.0


class TestVRAMEstimation:
    def test_4bit_fits_8gb(self):
        config = TrainingConfig(
            quantization="4bit",
            batch_size=2,
            max_seq_length=512,
            lora_enabled=True,
            lora_rank=16,
        )
        estimate = estimate_vram(config, vram_gb=8.0)
        assert estimate.total_gb > 0
        assert isinstance(estimate.fits_in_vram, bool)

    def test_no_quant_needs_more_vram(self):
        config_4bit = TrainingConfig(quantization="4bit", batch_size=4)
        config_none = TrainingConfig(quantization="none", batch_size=4)

        est_4bit = estimate_vram(config_4bit, vram_gb=24.0)
        est_none = estimate_vram(config_none, vram_gb=24.0)

        assert est_none.model_weights_gb > est_4bit.model_weights_gb

    def test_gradient_checkpointing_reduces_memory(self):
        config_no_gc = TrainingConfig(
            quantization="4bit",
            batch_size=4,
            use_gradient_checkpointing=False,
        )
        config_gc = TrainingConfig(
            quantization="4bit",
            batch_size=4,
            use_gradient_checkpointing=True,
        )
        est_no_gc = estimate_vram(config_no_gc, vram_gb=8.0)
        est_gc = estimate_vram(config_gc, vram_gb=8.0)

        assert est_gc.activations_gb < est_no_gc.activations_gb

    def test_larger_batch_uses_more_memory(self):
        config_small = TrainingConfig(quantization="4bit", batch_size=1)
        config_large = TrainingConfig(quantization="4bit", batch_size=8)

        est_small = estimate_vram(config_small, vram_gb=24.0)
        est_large = estimate_vram(config_large, vram_gb=24.0)

        assert est_large.total_gb > est_small.total_gb
