import numpy as np
import time
from src.data import generate_clustered_data, generate_ternary_weights
from src.engine import run_oxta_ternary, run_numpy_baseline

def test_ternary_logic():
    print("🧪 Testando Ternary Engine (Correção Matemática)...")
    N = 128
    block_size = 32

    # 1. Dados Pequenos
    A = np.random.rand(N, N).astype(np.float32)
    B_ternary = generate_ternary_weights(N, zero_prob=0.5)

    # 2. Baseline Real (Matmul padrão, mas convertendo B para float para o numpy aceitar)
    B_float = B_ternary.astype(np.float32)
    expected = A @ B_float

    # 3. Oxta Ternary
    # Precisamos criar active_blocks dummy (tudo ativo para teste de correção)
    n_blocks = N // block_size
    active_blocks = []
    for i in range(n_blocks):
        for k in range(n_blocks):
            active_blocks.append([i, k])
    active_blocks = np.array(active_blocks)

    # run_oxta_ternary returns time, but we need C_out?
    # Actually run_oxta_ternary allocates C internally and returns time.
    # We need to expose C for testing.
    # Or modify run_oxta_ternary to return C?
    # The original v2.2 signature returns time.
    # Let's import the kernel directly again for testing?
    # Or update run_oxta_ternary to return (time, C) or just C?
    # No, benchmark expects time.

    # Hack for test: Copy code or import kernel if possible.
    # I'll import the mask kernel.
    from src.engine import ternary_kernel_mask

    mask = np.ones((n_blocks, n_blocks), dtype=np.bool_)
    C_out = np.zeros((N, N), dtype=np.float32)
    ternary_kernel_mask(A, B_ternary, C_out, mask, block_size, N)

    # 4. Verificação
    # Pequena tolerância devido a ponto flutuante, mas deve ser muito próximo
    # Note que adição repetida pode ter erros de arredondamento diferentes de FMA.
    diff = np.abs(expected - C_out).max()
    print(f"Max Diff: {diff}")

    if diff < 1e-4:
        print("✅ Ternary Engine Aprovado!")
    else:
        print("❌ Falha na precisão matemática.")
        print("Expected Sample:", expected[0,0])
        print("Got Sample:", C_out[0,0])

if __name__ == "__main__":
    test_ternary_logic()
