import numpy as np
from src.data import generate_structured_2_4

def test_structured_sparsity():
    print("🧪 Testando Gerador de Esparsidade Estruturada 2:4...")
    N = 128
    B_float, B_ternary = generate_structured_2_4(N)

    # Check shape
    assert B_float.shape == (N, N)

    # Check logic: reshape to (N//4, 4, N) and count non-zeros along axis 1
    view = B_float.reshape(N//4, 4, N)
    nnz_counts = np.count_nonzero(view, axis=1)

    # All counts must be exactly 2
    if np.all(nnz_counts == 2):
        print("✅ Padrão 2:4 (Float) validado com sucesso!")
    else:
        print("❌ Falha na validação 2:4 (Float).")
        print("Max NNZ:", nnz_counts.max())
        print("Min NNZ:", nnz_counts.min())
        return

    # Check Ternary
    view_t = B_ternary.reshape(N//4, 4, N)
    nnz_counts_t = np.count_nonzero(view_t, axis=1)
    if np.all(nnz_counts_t == 2):
        print("✅ Padrão 2:4 (Ternary) validado com sucesso!")
        # Check values are -1, 0, 1
        unique_vals = np.unique(B_ternary)
        print(f"   Valores únicos encontrados: {unique_vals}")
    else:
        print("❌ Falha na validação 2:4 (Ternary).")

if __name__ == "__main__":
    test_structured_sparsity()
