import numpy as np
from src.bitpacking import pack_ternary_weights, unpack_ternary_weights

def test_bitpacking():
    print("🧪 Testando Bit-Packing (4:1 Compression)...")
    N = 128
    # Gera pesos ternários
    weights = np.random.choice([-1, 0, 1], size=(N, N)).astype(np.int8)

    # Pack
    packed = pack_ternary_weights(weights)

    # Check shape
    expected_shape = (N, N // 4)
    if packed.shape != expected_shape:
        print(f"❌ Erro de Shape: Esperado {expected_shape}, Recebido {packed.shape}")
        return

    # Unpack e comparar
    unpacked = unpack_ternary_weights(packed, (N, N))

    if np.array_equal(weights, unpacked):
        print("✅ Bit-Packing validado! Compressão 4x confirmada.")
        print(f"   Original Size: {weights.nbytes} bytes")
        print(f"   Packed Size:   {packed.nbytes} bytes")
    else:
        print("❌ Erro de Integridade dos Dados.")
        # Debug
        diff = weights != unpacked
        print("Indices diferentes:", np.argwhere(diff)[:5])
        print("Original:", weights[diff][:5])
        print("Unpacked:", unpacked[diff][:5])

if __name__ == "__main__":
    test_bitpacking()
