import torch

from orbitquant_wan_a2.nibbles import unpack_uint4
from orbitquant_wan_a2.source_quant import TARGET_SUFFIXES, is_target_key, _quantize_rows_to_packed


def test_exact_480_target_rule():
    keys = [f"blocks.{b}{s}" for b in range(40) for s in TARGET_SUFFIXES]
    assert len(keys) == 480
    assert len(set(keys)) == 480
    assert all(is_target_key(k) for k in keys)
    assert not is_target_key("blocks.0.block.modulation")
    assert not is_target_key("patch_embedding.weight")


def test_direct_quant_packs_codes_and_bf16_scale():
    torch.manual_seed(123)
    d = 8
    cb = torch.linspace(-0.8, 0.8, 16, dtype=torch.float32)
    bank_item = {
        "perm": torch.arange(d, dtype=torch.int64),
        "signs": torch.ones(d, dtype=torch.int8),
        "codebook": cb,
        "block_size": torch.tensor([8], dtype=torch.int32),
    }
    w = torch.randn(3, d, dtype=torch.bfloat16)
    packed, scale, stats = _quantize_rows_to_packed(w, bank_item, device=torch.device("cpu"), row_chunk=2)
    assert packed.dtype == torch.uint8
    assert packed.shape == (3, d // 2)
    assert scale.dtype == torch.bfloat16
    assert scale.shape == (3,)
    codes = unpack_uint4(packed, d)
    assert int(codes.min()) >= 0 and int(codes.max()) <= 15
    assert stats["row_scale_dtype"] == "bfloat16"


def test_safetensors_shape_probe_uses_slice_metadata(tmp_path):
    import json
    from safetensors.torch import save_file
    from orbitquant_wan_a2.source_quant import _shape_of

    shard = tmp_path / 'model-00001-of-00001.safetensors'
    save_file({'blocks.0.block.self_attn.q.weight': torch.zeros(5, 8, dtype=torch.bfloat16)}, str(shard))
    wm = {'blocks.0.block.self_attn.q.weight': shard.name}
    assert _shape_of(tmp_path, wm, 'blocks.0.block.self_attn.q.weight') == (5, 8)