from __future__ import annotations

import json

from safetensors.torch import save_file
import torch

from obliteratus.model_profile import default_self_improve_params, estimate_total_params, profile_model


def test_estimate_total_params_uses_nested_text_config():
    cfg = {
        "model_type": "qwen3_5",
        "text_config": {
            "model_type": "qwen3_5_text",
            "hidden_size": 128,
            "num_hidden_layers": 2,
            "num_attention_heads": 4,
            "num_key_value_heads": 2,
            "head_dim": 32,
            "intermediate_size": 256,
            "vocab_size": 1000,
        },
    }
    assert estimate_total_params(cfg) is not None
    assert estimate_total_params(cfg) > 0


def test_profile_model_counts_local_safetensors_exactly(tmp_path):
    model_dir = tmp_path / "model"
    model_dir.mkdir()
    (model_dir / "config.json").write_text(json.dumps({
        "model_type": "toy",
        "hidden_size": 8,
        "num_hidden_layers": 1,
        "intermediate_size": 16,
        "vocab_size": 32,
    }))
    save_file({
        "a.weight": torch.zeros(2, 3),
        "b.weight": torch.zeros(5),
    }, model_dir / "model.safetensors")

    profile = profile_model(str(model_dir), dtype="bfloat16")

    assert profile.source == "local_safetensors"
    assert profile.total_params == 11
    assert profile.total_params_b == 0.000000
    assert profile.hidden_size == 8


def test_default_self_improve_params_are_size_aware():
    from obliteratus.model_profile import ModelProfile
    big = ModelProfile("big", "test", int(27e9), 27.0, 27.0, 64, 5120, 17408, 248320, "qwen", "bf16")
    tiny = ModelProfile("tiny", "test", int(1e9), 1.0, 1.0, 16, 2048, 8192, 32000, "toy", "bf16")

    assert default_self_improve_params(big)["residue_weight"] < default_self_improve_params(tiny)["residue_weight"]
    assert default_self_improve_params(big)["n_directions"] >= default_self_improve_params(tiny)["n_directions"]