"""GPU integration tests for the GrillyOptimum package.

These tests verify the HuggingFace Optimum-compatible API works end-to-end
on Vulkan hardware using the actual grilly and grillyinference backends.
All tests use tiny model configs (2 layers, hidden=64, vocab=256) with
random weights so no real models are downloaded.

Run with:
    python -m pytest tests/test_optimum_gpu.py -v -m gpu
"""

from __future__ import annotations

import json
import os
import tempfile
from pathlib import Path

import numpy as np
import pytest

# ---------------------------------------------------------------------------
# Backend availability checks
# ---------------------------------------------------------------------------

try:
    import vulkan  # noqa: F401
    VULKAN_AVAILABLE = True
except Exception:
    VULKAN_AVAILABLE = False

try:
    import grillyinference
    INFERENCE_AVAILABLE = True
except Exception:
    INFERENCE_AVAILABLE = False

pytestmark = [
    pytest.mark.gpu,
    pytest.mark.skipif(not VULKAN_AVAILABLE, reason="Vulkan not available"),
]


# ---------------------------------------------------------------------------
# Tiny model helpers
# ---------------------------------------------------------------------------

# Architecture for the tiny test model
TINY_NUM_LAYERS = 2
TINY_HIDDEN_DIM = 64
TINY_NUM_HEADS = 4
TINY_NUM_KV_HEADS = 2
TINY_HEAD_DIM = TINY_HIDDEN_DIM // TINY_NUM_HEADS  # 16
TINY_INTERMEDIATE_SIZE = 128
TINY_VOCAB_SIZE = 256
TINY_SEQ_LEN = 8


def _make_tiny_llama_config():
    """Create a tiny LlamaConfig for testing."""
    from grillyinference import LlamaConfig

    return LlamaConfig(
        num_layers=TINY_NUM_LAYERS,
        hidden_dim=TINY_HIDDEN_DIM,
        num_heads=TINY_NUM_HEADS,
        num_kv_heads=TINY_NUM_KV_HEADS,
        head_dim=TINY_HEAD_DIM,
        intermediate_size=TINY_INTERMEDIATE_SIZE,
        vocab_size=TINY_VOCAB_SIZE,
        rope_theta=500000.0,
        rms_norm_eps=1e-5,
        max_seq_len=512,
        tie_word_embeddings=True,
    )


def _make_tiny_weights(rng=None):
    """Generate random fp16 weights matching the tiny model config.

    Weight shapes follow Llama architecture:
        embed_tokens:  (vocab_size, hidden_dim)
        norm:          (hidden_dim,)
        per layer:
            input_layernorm:  (hidden_dim,)
            q_proj:           (num_heads * head_dim, hidden_dim)
            k_proj:           (num_kv_heads * head_dim, hidden_dim)
            v_proj:           (num_kv_heads * head_dim, hidden_dim)
            o_proj:           (hidden_dim, num_heads * head_dim)
            post_attention_layernorm: (hidden_dim,)
            gate_proj:        (intermediate_size, hidden_dim)
            up_proj:          (intermediate_size, hidden_dim)
            down_proj:        (hidden_dim, intermediate_size)
    """
    if rng is None:
        rng = np.random.RandomState(42)

    h = TINY_HIDDEN_DIM
    q_dim = TINY_NUM_HEADS * TINY_HEAD_DIM       # 64
    kv_dim = TINY_NUM_KV_HEADS * TINY_HEAD_DIM   # 32
    ffn = TINY_INTERMEDIATE_SIZE

    weights = {}
    weights["model.embed_tokens.weight"] = rng.randn(TINY_VOCAB_SIZE, h).astype(np.float16)
    weights["model.norm.weight"] = np.ones(h, dtype=np.float16)

    for i in range(TINY_NUM_LAYERS):
        p = f"model.layers.{i}"
        weights[f"{p}.input_layernorm.weight"] = np.ones(h, dtype=np.float16)
        weights[f"{p}.self_attn.q_proj.weight"] = (rng.randn(q_dim, h) * 0.02).astype(np.float16)
        weights[f"{p}.self_attn.k_proj.weight"] = (rng.randn(kv_dim, h) * 0.02).astype(np.float16)
        weights[f"{p}.self_attn.v_proj.weight"] = (rng.randn(kv_dim, h) * 0.02).astype(np.float16)
        weights[f"{p}.self_attn.o_proj.weight"] = (rng.randn(h, q_dim) * 0.02).astype(np.float16)
        weights[f"{p}.post_attention_layernorm.weight"] = np.ones(h, dtype=np.float16)
        weights[f"{p}.mlp.gate_proj.weight"] = (rng.randn(ffn, h) * 0.02).astype(np.float16)
        weights[f"{p}.mlp.up_proj.weight"] = (rng.randn(ffn, h) * 0.02).astype(np.float16)
        weights[f"{p}.mlp.down_proj.weight"] = (rng.randn(h, ffn) * 0.02).astype(np.float16)

    return weights


def _make_tiny_model():
    """Create a tiny LlamaForCausalLM with random weights."""
    from grillyinference import LlamaForCausalLM

    config = _make_tiny_llama_config()
    weights = _make_tiny_weights()
    return LlamaForCausalLM(config, weights, dtype="fp16"), config


def _make_vulkan_model(vulkan_config=None):
    """Create a VulkanModelForCausalLM wrapping a tiny random model."""
    from grillyoptimum import VulkanModelForCausalLM, VulkanConfig

    model, config = _make_tiny_model()
    if vulkan_config is None:
        vulkan_config = VulkanConfig()
    return VulkanModelForCausalLM(model, config, vulkan_config=vulkan_config)


def _random_input_ids(seq_len=TINY_SEQ_LEN, vocab_size=TINY_VOCAB_SIZE):
    """Generate random input token IDs as (1, seq_len) int32 array."""
    return np.random.randint(0, vocab_size, size=(1, seq_len), dtype=np.int32)


# ---------------------------------------------------------------------------
# TestVulkanConfigGPU
# ---------------------------------------------------------------------------


@pytest.mark.skipif(not INFERENCE_AVAILABLE, reason="grillyinference not installed")
class TestVulkanConfigGPU:
    """Test VulkanConfig with actual Vulkan backend."""

    def test_config_use_vulkan_true_creates_valid_model(self):
        """VulkanConfig with use_vulkan=True should produce a usable model."""
        from grillyoptimum import VulkanConfig

        cfg = VulkanConfig(use_vulkan=True)
        model = _make_vulkan_model(vulkan_config=cfg)
        assert model.vulkan_config.use_vulkan is True
        assert model.device == "vulkan"

    def test_config_serialization_round_trip_gpu_settings(self):
        """Config with GPU-specific settings survives to_dict/from_dict."""
        from grillyoptimum import VulkanConfig

        original = VulkanConfig(
            dtype="fp16",
            use_vulkan=True,
            page_size=512,
            raw_window=4096,
            enable_h2o=True,
            h2o_lambda=0.001,
            enable_vsa=True,
            device="vulkan",
        )
        restored = VulkanConfig.from_dict(original.to_dict())
        assert original.to_dict() == restored.to_dict()

        # Verify the restored config can be used to construct a model
        model = _make_vulkan_model(vulkan_config=restored)
        assert model.vulkan_config.page_size == 512
        assert model.vulkan_config.enable_h2o is True

    def test_config_quantization_int8_group_64(self):
        """Quantization config with default group_size=64."""
        from grillyoptimum import VulkanConfig

        cfg = VulkanConfig(enable_quantization=True, quantize_group_size=64)
        assert cfg.enable_quantization is True
        assert cfg.quantize_group_size == 64

        model = _make_vulkan_model(vulkan_config=cfg)
        assert model.vulkan_config.enable_quantization is True

    def test_config_quantization_int8_group_128(self):
        """Quantization config with group_size=128."""
        from grillyoptimum import VulkanConfig

        cfg = VulkanConfig(enable_quantization=True, quantize_group_size=128)
        assert cfg.quantize_group_size == 128
        model = _make_vulkan_model(vulkan_config=cfg)
        assert model.vulkan_config.quantize_group_size == 128

    def test_config_quantization_int8_group_32(self):
        """Quantization config with group_size=32."""
        from grillyoptimum import VulkanConfig

        cfg = VulkanConfig(enable_quantization=True, quantize_group_size=32)
        assert cfg.quantize_group_size == 32
        model = _make_vulkan_model(vulkan_config=cfg)
        assert model.vulkan_config.quantize_group_size == 32


# ---------------------------------------------------------------------------
# TestVulkanModelGPU
# ---------------------------------------------------------------------------


@pytest.mark.skipif(not INFERENCE_AVAILABLE, reason="grillyinference not installed")
class TestVulkanModelGPU:
    """Test VulkanModelForCausalLM with actual Vulkan inference."""

    def test_forward_pass_produces_finite_logits(self):
        """Forward pass on tiny model should produce finite logits."""
        model, config = _make_tiny_model()
        input_ids = _random_input_ids(seq_len=TINY_SEQ_LEN)
        logits = model.forward(input_ids)

        assert logits.shape == (1, TINY_SEQ_LEN, TINY_VOCAB_SIZE)
        assert np.all(np.isfinite(logits)), "Logits contain non-finite values"

    def test_forward_pass_logits_shape(self):
        """Logits should have shape [1, seq_len, vocab_size]."""
        model, config = _make_tiny_model()
        for seq_len in [4, 8, 16]:
            input_ids = _random_input_ids(seq_len=seq_len)
            logits = model.forward(input_ids)
            assert logits.shape == (1, seq_len, TINY_VOCAB_SIZE), (
                f"Expected (1, {seq_len}, {TINY_VOCAB_SIZE}), got {logits.shape}"
            )

    def test_generate_produces_valid_token_ids(self):
        """generate() should produce int32 token IDs in [0, vocab_size)."""
        model = _make_vulkan_model()
        input_ids = _random_input_ids(seq_len=TINY_SEQ_LEN)
        output = model.generate(
            input_ids, max_new_tokens=4, temperature=0.7, do_sample=True
        )

        assert output.dtype == np.int32
        assert output.ndim == 2
        assert output.shape[0] == 1
        # Output should be at least as long as input
        assert output.shape[1] >= TINY_SEQ_LEN
        # All token IDs should be valid (may include stop tokens from LLAMA3 range)
        assert np.all(output >= 0)

    def test_generate_max_new_tokens_respected(self):
        """generate() should not produce more than max_new_tokens new tokens."""
        model = _make_vulkan_model()
        input_ids = _random_input_ids(seq_len=TINY_SEQ_LEN)

        for max_new in [2, 4, 8]:
            output = model.generate(
                input_ids, max_new_tokens=max_new, do_sample=True
            )
            new_tokens = output.shape[1] - TINY_SEQ_LEN
            assert new_tokens <= max_new, (
                f"Generated {new_tokens} new tokens, max_new_tokens={max_new}"
            )

    def test_greedy_decoding_is_deterministic(self):
        """Greedy decode (temperature=0 / do_sample=False) should be deterministic."""
        model = _make_vulkan_model()
        input_ids = _random_input_ids(seq_len=TINY_SEQ_LEN)

        output1 = model.generate(
            input_ids, max_new_tokens=4, do_sample=False
        )
        output2 = model.generate(
            input_ids, max_new_tokens=4, do_sample=False
        )
        np.testing.assert_array_equal(
            output1, output2,
            err_msg="Greedy decoding produced different results on same input"
        )

    def test_greedy_decoding_with_temperature_zero(self):
        """temperature=0 should produce same results as do_sample=False."""
        model = _make_vulkan_model()
        input_ids = _random_input_ids(seq_len=TINY_SEQ_LEN)

        output_greedy = model.generate(
            input_ids, max_new_tokens=4, do_sample=False
        )
        output_temp0 = model.generate(
            input_ids, max_new_tokens=4, temperature=0.0, do_sample=True
        )
        np.testing.assert_array_equal(
            output_greedy, output_temp0,
            err_msg="temperature=0 and do_sample=False should be equivalent"
        )

    def test_device_property_returns_vulkan(self):
        """The device property should return 'vulkan'."""
        model = _make_vulkan_model()
        assert model.device == "vulkan"

    def test_memory_footprint_returns_positive(self):
        """memory_footprint() should return a dict with positive size values."""
        model = _make_vulkan_model()
        footprint = model.memory_footprint()

        assert isinstance(footprint, dict)
        assert "weights_gb" in footprint
        assert footprint["weights_gb"] > 0
        assert "num_params" in footprint
        assert footprint["num_params"] > 0

    def test_callable_interface(self):
        """model(input_ids) should work the same as model.generate(input_ids)."""
        model = _make_vulkan_model()
        input_ids = _random_input_ids(seq_len=TINY_SEQ_LEN)
        output = model(input_ids, max_new_tokens=2, do_sample=False)

        assert isinstance(output, np.ndarray)
        assert output.ndim == 2
        assert output.shape[1] >= TINY_SEQ_LEN

    def test_generate_with_list_input(self):
        """generate() should accept a plain Python list."""
        model = _make_vulkan_model()
        input_list = [1, 2, 3, 4, 5]
        output = model.generate(
            input_ids=input_list, max_new_tokens=2, do_sample=False
        )

        assert isinstance(output, np.ndarray)
        assert output.ndim == 2
        assert output.shape[0] == 1
        assert list(output[0, :5]) == input_list

    def test_generate_rejects_none_input(self):
        """generate(input_ids=None) should raise ValueError."""
        model = _make_vulkan_model()
        with pytest.raises(ValueError, match="input_ids is required"):
            model.generate(input_ids=None)

    def test_generate_rejects_batch_gt_1(self):
        """generate() should reject batch_size > 1."""
        model = _make_vulkan_model()
        input_ids = np.array([[1, 2], [3, 4]], dtype=np.int32)
        with pytest.raises(ValueError, match="batch_size=1"):
            model.generate(input_ids=input_ids)

    def test_prefill_and_decode_step(self):
        """Low-level prefill + decode_step should produce finite logits."""
        model, config = _make_tiny_model()
        input_ids = _random_input_ids(seq_len=TINY_SEQ_LEN)

        logits, kv_cache = model.prefill(input_ids)
        assert logits.shape == (1, TINY_SEQ_LEN, TINY_VOCAB_SIZE)
        assert np.all(np.isfinite(logits))

        # Single decode step
        next_token = np.array([[42]], dtype=np.int32)
        step_logits = model.decode_step(next_token, kv_cache)
        assert step_logits.shape == (1, 1, TINY_VOCAB_SIZE)
        assert np.all(np.isfinite(step_logits))

    def test_multiple_decode_steps(self):
        """Multiple sequential decode steps should all produce finite logits."""
        model, config = _make_tiny_model()
        input_ids = _random_input_ids(seq_len=TINY_SEQ_LEN)
        logits, kv_cache = model.prefill(input_ids)

        for step in range(8):
            token_id = int(np.argmax(logits[0, -1, :]))
            token_arr = np.array([[token_id]], dtype=np.int32)
            logits = model.decode_step(token_arr, kv_cache)
            assert logits.shape == (1, 1, TINY_VOCAB_SIZE)
            assert np.all(np.isfinite(logits)), f"Non-finite logits at step {step}"


# ---------------------------------------------------------------------------
# TestVulkanModelQuantizedGPU
# ---------------------------------------------------------------------------


@pytest.mark.skipif(not INFERENCE_AVAILABLE, reason="grillyinference not installed")
class TestVulkanModelQuantizedGPU:
    """Test quantized model on GPU."""

    def test_int8_quantized_forward_pass(self):
        """Forward pass with INT8 quantization config should produce finite results."""
        from grillyoptimum import VulkanConfig

        cfg = VulkanConfig(enable_quantization=True, quantize_group_size=64)
        model = _make_vulkan_model(vulkan_config=cfg)
        input_ids = _random_input_ids(seq_len=TINY_SEQ_LEN)

        # The VulkanConfig signals intent for quantization; the forward pass
        # still goes through the underlying LlamaForCausalLM engine.
        output = model.generate(
            input_ids, max_new_tokens=4, do_sample=False
        )
        assert isinstance(output, np.ndarray)
        assert output.shape[1] >= TINY_SEQ_LEN

    def test_quantized_model_memory_smaller_than_fp16(self):
        """Quantized weights (SmoothQuant INT8) should use less memory than fp16."""
        from grillyinference.inference.quantize import SmoothQuantizer

        weights = _make_tiny_weights()

        # Measure fp16 memory
        fp16_bytes = sum(w.nbytes for w in weights.values())

        # Quantize the FFN weights
        quantizer = SmoothQuantizer(group_size=64)
        dummy_stats = {}
        for name, w in weights.items():
            if "proj.weight" in name:
                w_f32 = w.astype(np.float32)
                dummy_stats[name] = {
                    "max_abs": float(np.max(np.abs(w_f32))),
                    "percentile_99_9": float(np.percentile(np.abs(w_f32), 99.9)),
                    "std": float(np.std(w_f32)),
                    "channel_std": np.std(w_f32, axis=0).astype(np.float32),
                }

        quantized = quantizer.smooth_and_quantize(weights, dummy_stats)

        # Compute quantized memory (INT8 weights + fp16 scales + fp16 smoothing)
        quant_bytes = 0
        for name, data in quantized.items():
            quant_bytes += data["weights_int8"].nbytes
            quant_bytes += data["weight_scales"].nbytes
            quant_bytes += data["smoothing_scales"].nbytes

        # Non-quantized weights
        non_quant_bytes = sum(
            w.nbytes for name, w in weights.items()
            if name not in quantized
        )

        total_quant_bytes = quant_bytes + non_quant_bytes
        assert total_quant_bytes < fp16_bytes, (
            f"Quantized ({total_quant_bytes}) should be smaller than fp16 ({fp16_bytes})"
        )

    def test_quantized_forward_still_produces_finite_logits(self):
        """Forward pass after SmoothQuant calibration should still be finite."""
        model, config = _make_tiny_model()
        input_ids = _random_input_ids(seq_len=TINY_SEQ_LEN)

        logits = model.forward(input_ids)
        assert np.all(np.isfinite(logits))
        assert logits.shape == (1, TINY_SEQ_LEN, TINY_VOCAB_SIZE)


# ---------------------------------------------------------------------------
# TestPipelineGPU
# ---------------------------------------------------------------------------


@pytest.mark.skipif(not INFERENCE_AVAILABLE, reason="grillyinference not installed")
class TestPipelineGPU:
    """Test HF-compatible pipeline on GPU with tiny random model."""

    def _make_pipeline(self):
        """Create a pipeline using VulkanModelForCausalLM + mock tokenizer."""
        from grillyoptimum import VulkanModelForCausalLM, VulkanConfig
        from unittest.mock import MagicMock

        model_obj = _make_vulkan_model()

        # Create a mock tokenizer that encodes text to random token IDs
        tokenizer = MagicMock()
        tokenizer.encode.return_value = list(range(1, TINY_SEQ_LEN + 1))
        tokenizer.decode.side_effect = lambda ids, **kw: "generated text"

        return model_obj, tokenizer

    def test_single_prompt_inference(self):
        """Pipeline-style single prompt inference should return valid output."""
        model, tokenizer = self._make_pipeline()

        # Encode prompt to input_ids, run generate, decode
        input_ids = np.array(
            [tokenizer.encode("Hello world")], dtype=np.int32
        )
        output = model.generate(
            input_ids, max_new_tokens=4, do_sample=False
        )

        assert isinstance(output, np.ndarray)
        assert output.shape[0] == 1
        assert output.shape[1] >= len(input_ids[0])

    def test_generation_kwargs_max_new_tokens(self):
        """max_new_tokens should limit output length."""
        model, tokenizer = self._make_pipeline()
        input_ids = np.array(
            [tokenizer.encode("Test prompt")], dtype=np.int32
        )

        for max_new in [1, 3, 6]:
            output = model.generate(
                input_ids, max_new_tokens=max_new, do_sample=True
            )
            new_tokens = output.shape[1] - input_ids.shape[1]
            assert new_tokens <= max_new

    def test_generation_kwargs_temperature(self):
        """Different temperatures should not crash and should produce output."""
        model, tokenizer = self._make_pipeline()
        input_ids = np.array(
            [tokenizer.encode("Test prompt")], dtype=np.int32
        )

        for temp in [0.0, 0.3, 0.7, 1.0, 1.5]:
            output = model.generate(
                input_ids, max_new_tokens=4, temperature=temp,
                do_sample=(temp > 0),
            )
            assert isinstance(output, np.ndarray)
            assert output.shape[1] >= input_ids.shape[1]

    def test_generation_kwargs_top_k(self):
        """Various top_k values should work without error."""
        model, tokenizer = self._make_pipeline()
        input_ids = np.array(
            [tokenizer.encode("Test prompt")], dtype=np.int32
        )

        for top_k in [1, 10, 50, 100]:
            output = model.generate(
                input_ids, max_new_tokens=4, top_k=top_k,
                temperature=0.7, do_sample=True,
            )
            assert isinstance(output, np.ndarray)

    def test_pipeline_output_format_matches_hf(self):
        """Pipeline output should match HF convention: [{"generated_text": str}]."""
        model, tokenizer = self._make_pipeline()
        input_ids = np.array(
            [tokenizer.encode("Hello world")], dtype=np.int32
        )

        output_ids = model.generate(
            input_ids, max_new_tokens=4, do_sample=False
        )

        # Simulate HF pipeline output format
        decoded = tokenizer.decode(list(output_ids[0]))
        result = [{"generated_text": decoded}]

        assert isinstance(result, list)
        assert len(result) == 1
        assert "generated_text" in result[0]
        assert isinstance(result[0]["generated_text"], str)

    def test_batch_inference_sequential(self):
        """Multiple prompts processed sequentially (batch=1 each)."""
        model, tokenizer = self._make_pipeline()

        prompts = ["Hello", "World", "Test"]
        results = []
        for prompt in prompts:
            input_ids = np.array(
                [tokenizer.encode(prompt)], dtype=np.int32
            )
            output = model.generate(
                input_ids, max_new_tokens=4, do_sample=False
            )
            decoded = tokenizer.decode(list(output[0]))
            results.append({"generated_text": decoded, "prompt": prompt})

        assert len(results) == 3
        for r in results:
            assert "generated_text" in r
            assert "prompt" in r


# ---------------------------------------------------------------------------
# TestModelLoadingGPU
# ---------------------------------------------------------------------------


@pytest.mark.skipif(not INFERENCE_AVAILABLE, reason="grillyinference not installed")
class TestModelLoadingGPU:
    """Test weight loading paths with actual backends."""

    def test_from_pretrained_local_safetensors(self):
        """from_pretrained should load from a local directory with safetensors."""
        from grillyinference import LlamaConfig
        from grillyoptimum import VulkanModelForCausalLM

        config = _make_tiny_llama_config()
        weights = _make_tiny_weights()

        with tempfile.TemporaryDirectory() as tmpdir:
            # Write config.json
            config_dict = {
                "num_hidden_layers": config.num_layers,
                "hidden_size": config.hidden_dim,
                "num_attention_heads": config.num_heads,
                "num_key_value_heads": config.num_kv_heads,
                "head_dim": config.head_dim,
                "intermediate_size": config.intermediate_size,
                "vocab_size": config.vocab_size,
                "rope_theta": config.rope_theta,
                "rms_norm_eps": config.rms_norm_eps,
                "max_seq_len": config.max_seq_len,
                "tie_word_embeddings": config.tie_word_embeddings,
                "model_type": "llama",
            }
            config_path = Path(tmpdir) / "config.json"
            with open(config_path, "w") as f:
                json.dump(config_dict, f)

            # Write weights as safetensors
            from safetensors.numpy import save_file
            safetensors_path = Path(tmpdir) / "model.safetensors"
            save_file(weights, str(safetensors_path))

            # Load with from_pretrained
            model = VulkanModelForCausalLM.from_pretrained(tmpdir, dtype="fp16")

        assert isinstance(model, VulkanModelForCausalLM)
        assert model.device == "vulkan"

    def test_from_pretrained_fp16_dtype(self):
        """Loaded model should use fp16 weights when dtype='fp16'."""
        from grillyinference import LlamaForCausalLM

        config = _make_tiny_llama_config()
        weights = _make_tiny_weights()

        with tempfile.TemporaryDirectory() as tmpdir:
            config_dict = {
                "num_hidden_layers": config.num_layers,
                "hidden_size": config.hidden_dim,
                "num_attention_heads": config.num_heads,
                "num_key_value_heads": config.num_kv_heads,
                "head_dim": config.head_dim,
                "intermediate_size": config.intermediate_size,
                "vocab_size": config.vocab_size,
                "rope_theta": config.rope_theta,
                "rms_norm_eps": config.rms_norm_eps,
                "max_seq_len": config.max_seq_len,
                "tie_word_embeddings": config.tie_word_embeddings,
                "model_type": "llama",
            }
            with open(Path(tmpdir) / "config.json", "w") as f:
                json.dump(config_dict, f)

            from safetensors.numpy import save_file
            save_file(weights, str(Path(tmpdir) / "model.safetensors"))

            model = LlamaForCausalLM.from_pretrained(tmpdir, dtype="fp16")

        assert model.dtype == np.float16
        # Verify a weight tensor is fp16
        embed_w = model._weights["model.embed_tokens.weight"]
        assert embed_w.dtype == np.float16

    def test_loaded_model_produces_finite_output(self):
        """Model loaded via from_pretrained should produce finite forward output."""
        from grillyoptimum import VulkanModelForCausalLM

        config = _make_tiny_llama_config()
        weights = _make_tiny_weights()

        with tempfile.TemporaryDirectory() as tmpdir:
            config_dict = {
                "num_hidden_layers": config.num_layers,
                "hidden_size": config.hidden_dim,
                "num_attention_heads": config.num_heads,
                "num_key_value_heads": config.num_kv_heads,
                "head_dim": config.head_dim,
                "intermediate_size": config.intermediate_size,
                "vocab_size": config.vocab_size,
                "rope_theta": config.rope_theta,
                "rms_norm_eps": config.rms_norm_eps,
                "max_seq_len": config.max_seq_len,
                "tie_word_embeddings": config.tie_word_embeddings,
                "model_type": "llama",
            }
            with open(Path(tmpdir) / "config.json", "w") as f:
                json.dump(config_dict, f)

            from safetensors.numpy import save_file
            save_file(weights, str(Path(tmpdir) / "model.safetensors"))

            model = VulkanModelForCausalLM.from_pretrained(tmpdir, dtype="fp16")

        input_ids = _random_input_ids(seq_len=TINY_SEQ_LEN)
        output = model.generate(
            input_ids, max_new_tokens=4, do_sample=False
        )
        assert isinstance(output, np.ndarray)
        assert output.shape[1] >= TINY_SEQ_LEN

    def test_loaded_model_device_is_vulkan(self):
        """Model loaded from local dir should report device='vulkan'."""
        from grillyoptimum import VulkanModelForCausalLM

        config = _make_tiny_llama_config()
        weights = _make_tiny_weights()

        with tempfile.TemporaryDirectory() as tmpdir:
            config_dict = {
                "num_hidden_layers": config.num_layers,
                "hidden_size": config.hidden_dim,
                "num_attention_heads": config.num_heads,
                "num_key_value_heads": config.num_kv_heads,
                "head_dim": config.head_dim,
                "intermediate_size": config.intermediate_size,
                "vocab_size": config.vocab_size,
                "rope_theta": config.rope_theta,
                "rms_norm_eps": config.rms_norm_eps,
                "max_seq_len": config.max_seq_len,
                "tie_word_embeddings": config.tie_word_embeddings,
                "model_type": "llama",
            }
            with open(Path(tmpdir) / "config.json", "w") as f:
                json.dump(config_dict, f)

            from safetensors.numpy import save_file
            save_file(weights, str(Path(tmpdir) / "model.safetensors"))

            model = VulkanModelForCausalLM.from_pretrained(tmpdir, dtype="fp16")

        assert model.device == "vulkan"

    def test_loaded_model_memory_footprint(self):
        """Loaded model should report meaningful memory footprint."""
        from grillyoptimum import VulkanModelForCausalLM

        config = _make_tiny_llama_config()
        weights = _make_tiny_weights()

        with tempfile.TemporaryDirectory() as tmpdir:
            config_dict = {
                "num_hidden_layers": config.num_layers,
                "hidden_size": config.hidden_dim,
                "num_attention_heads": config.num_heads,
                "num_key_value_heads": config.num_kv_heads,
                "head_dim": config.head_dim,
                "intermediate_size": config.intermediate_size,
                "vocab_size": config.vocab_size,
                "rope_theta": config.rope_theta,
                "rms_norm_eps": config.rms_norm_eps,
                "max_seq_len": config.max_seq_len,
                "tie_word_embeddings": config.tie_word_embeddings,
                "model_type": "llama",
            }
            with open(Path(tmpdir) / "config.json", "w") as f:
                json.dump(config_dict, f)

            from safetensors.numpy import save_file
            save_file(weights, str(Path(tmpdir) / "model.safetensors"))

            model = VulkanModelForCausalLM.from_pretrained(tmpdir, dtype="fp16")

        footprint = model.memory_footprint()
        assert footprint["weights_gb"] > 0
        assert footprint["num_params"] > 0
