import numpy as np

# --- 1. GERADOR DE DADOS BIOMIMÉTICOS (CLUSTERED SPARSITY) ---
def generate_clustered_data(n, block_size, zero_prob):
    """
    Gera dados onde os zeros não são ruído branco, mas sim 'manchas' de inatividade,
    simulando o comportamento biológico ou ativações de redes neurais reais.
    """
    n_blocks = n // block_size
    # Máscara de baixa resolução (decisão por bloco)
    block_mask = np.random.rand(n_blocks, n_blocks) > zero_prob
    # Expande a máscara para a resolução total
    full_mask = np.kron(block_mask, np.ones((block_size, block_size)))
    data = np.random.rand(n, n).astype(np.float32)
    return data * full_mask

def generate_ternary_weights(n, zero_prob=0.0):
    """
    Gera pesos ternários {-1, 0, 1} armazenados como int8.
    Isso simula a compressão de memória (4x menor que float32).
    """
    # Gera números aleatórios e mapeia para -1, 0, 1
    # Ex: < 0.33 -> -1, > 0.66 -> 1, resto -> 0 (se zero_prob=0)
    # Se zero_prob > 0, ajustamos os limiares.

    data = np.random.choice([-1, 0, 1], size=(n, n), p=[(1-zero_prob)/2, zero_prob, (1-zero_prob)/2])
    return data.astype(np.int8)

def generate_structured_2_4(n):
    """
    Gera pesos com esparsidade estruturada 2:4 (NVIDIA Ampere Style).
    Ao longo da dimensão de entrada (linhas, se B for KxN e K é input channel),
    cada grupo de 4 elementos tem exatamente 2 zeros.
    Retorna matriz de float32 (para compatibilidade baseline) e int8 (ternário).
    """
    # 6 Permutações válidas para 2 zeros em 4 posições
    patterns = np.array([
        [0, 0, 1, 1],
        [0, 1, 0, 1],
        [0, 1, 1, 0],
        [1, 0, 0, 1],
        [1, 0, 1, 0],
        [1, 1, 0, 0]
    ], dtype=np.int8)

    # K=N. Assumimos N divisível por 4.
    assert n % 4 == 0, "N must be divisible by 4 for 2:4 sparsity"

    n_groups = n // 4

    # Escolhemos um padrão aleatório para cada grupo de 4 linhas, para cada coluna.
    # Shape: (N_groups, 1, N) -> Broadcasting patterns
    # Mas queremos variar por coluna também? Geralmente sim.
    # Shape alvo: (N, N). Reshape -> (N_groups, 4, N).

    choice_indices = np.random.randint(0, 6, size=(n_groups, n)) # Escolha por grupo e por coluna

    # Construindo a máscara
    # patterns[choice_indices] -> (n_groups, n, 4). Precisamos transpor o 4 para o meio?
    # O reshape de B será (n_groups, 4, n).
    # Então queremos que para cada (group, col) tenhamos um vetor de 4.

    mask_parts = patterns[choice_indices] # shape (n_groups, n, 4)
    # Transpose para (n_groups, 4, n) para combinar com reshape da matriz
    mask_parts = mask_parts.transpose(0, 2, 1)

    mask = mask_parts.reshape(n, n)

    # Preenche valores aleatórios
    data = np.random.rand(n, n).astype(np.float32) - 0.5
    data_sparse = data * mask

    # Versão Ternária: -1 ou 1 onde mask é 1
    ternary_vals = np.sign(data)
    ternary_vals[ternary_vals == 0] = 1 # fix zero sign
    ternary_sparse = (ternary_vals * mask).astype(np.int8)

    return data_sparse, ternary_sparse
