import numpy as np

class BlockPruner:
    """
    Simula o 'Sparsity-Aware Training' (The Teacher).
    Poda (Pruning) de modelos densos para o formato Block-Sparse do Maxwell.
    """
    def __init__(self, block_size=128):
        self.block_size = block_size

    def prune(self, weights, sparsity_target=0.8):
        """
        Recebe uma matriz de pesos Densa e força a esparsidade por blocos.
        Critério: L1 Norm do bloco.
        """
        n, m = weights.shape
        n_blocks_row = n // self.block_size
        n_blocks_col = m // self.block_size

        # 1. Calcular Energia dos Blocos (L1 Norm)
        # Reshape para (Nb_row, Block, Nb_col, Block)
        # Swap axes para (Nb_row, Nb_col, Block, Block)
        view = weights.reshape(n_blocks_row, self.block_size, n_blocks_col, self.block_size)
        view = view.transpose(0, 2, 1, 3)

        # Soma absoluta (Energia)
        block_energies = np.sum(np.abs(view), axis=(2, 3)) # Shape (Nb_row, Nb_col)

        # 2. Determinar Limiar (Threshold)
        # Queremos matar os X% blocos com menor energia
        flat_energies = block_energies.flatten()
        threshold = np.percentile(flat_energies, sparsity_target * 100)

        # 3. Criar Máscara
        # Blocos com energia > threshold sobrevivem
        active_mask = block_energies > threshold

        # 4. Aplicar Máscara (Expandir de volta)
        # Kronecker product para expandir a máscara para o tamanho original
        full_mask = np.kron(active_mask, np.ones((self.block_size, self.block_size)))

        pruned_weights = weights * full_mask

        # Stats
        real_sparsity = 1.0 - np.mean(full_mask)
        print(f"✂️  Pruning Completo. Esparsidade Alvo: {sparsity_target:.0%}, Real: {real_sparsity:.1%}")

        return pruned_weights.astype(np.float32)
