import time
import numpy as np
from numba import njit, prange
try:
    import maxwell_core
except ImportError:
    print("Warning: maxwell_core not found. C++ backends will fail.")
    maxwell_core = None

def run_maxwell_forward(A, B_packed, C, backend='cpu'):
    """
    Executes the Forward Pass using the Maxwell C++ Engine.
    backend: 'cpu' (AVX2 Tiled) or 'gpu' (CUDA CSR)
    """
    if maxwell_core is None:
        raise ImportError("maxwell_core module not compiled/loaded")
    
    start = time.time()
    if backend == 'gpu':
        # GPU CSR Persistent Kernel
        # Note: B_packed must be a PackedWeights object
        maxwell_core.run_maxwell_cuda_csr(A, B_packed, C)
    else:
        # CPU AVX2 Tiled Kernel
        maxwell_core.run_maxwell_metal(A, B_packed, C)
        
    return time.time() - start

def run_maxwell_backward(grad_C, B_packed, grad_A):
    """
    Executes the Backward Pass (Data Gradients) using Maxwell C++.
    Computes grad_A = grad_C * B.T
    """
    if maxwell_core is None:
        raise ImportError("maxwell_core module not compiled/loaded")
        
    start = time.time()
    maxwell_core.run_maxwell_backward(grad_C, B_packed, grad_A)
    return time.time() - start

# --- 3. MOTORES DE INFERÊNCIA ---

@njit(parallel=True, fastmath=True)
def ternary_kernel_mask(A, B, C, mask, block_size, N):
    """
    Kernel otimizado com paralelismo por linha de blocos (Thread-Safe).
    Usa máscara booleana para identificar blocos ativos.
    Elimina FPU MUL, usando apenas ADD/SUB.
    """
    n_b_rows, n_b_cols = mask.shape

    # Paralelização no eixo externo (Linhas de Blocos) garante que
    # threads diferentes escrevam em linhas diferentes de C. Sem Race Condition.
    for bi in prange(n_b_rows):
        i_start = bi * block_size

        for bk in range(n_b_cols):
            if not mask[bi, bk]:
                continue

            k_start = bk * block_size

            # Loop Tiling para o bloco atual
            for i_local in range(block_size):
                i_global = i_start + i_local

                # Carregar linha de A para registradores/cache
                # Acesso: A[i_global, k_start : k_start+block_size]

                for k_local in range(block_size):
                    k_global = k_start + k_local
                    val = A[i_global, k_global]

                    if val == 0: continue

                    # B assume valores int8 {-1, 0, 1}
                    # Acessamos a linha inteira de B.
                    # row_B = B[k_global, :]
                    # row_C = C[i_global, :]

                    # Vetorização interna (simulada pelo compilador Numba/LLVM)
                    for j in range(N):
                        w = B[k_global, j]
                        # Substituição de MUL por Lógica Ternária
                        if w == 1:
                            C[i_global, j] += val
                        elif w == -1:
                            C[i_global, j] -= val
                        # w == 0: do nothing

def run_oxta_ternary(A, B, active_blocks, block_size, n):
    """
    Wrapper para o motor Ternário.
    Converte lista de blocos para máscara para permitir paralelismo eficiente.
    """
    start = time.time()

    # Convert index list to boolean mask
    n_b = n // block_size
    mask = np.zeros((n_b, n_b), dtype=np.bool_)
    # active_blocks é (N_active, 2)
    if len(active_blocks) > 0:
        mask[active_blocks[:, 0], active_blocks[:, 1]] = True

    C = np.zeros((n, n), dtype=np.float32)
    ternary_kernel_mask(A, B, C, mask, block_size, n)
    return time.time() - start

def run_numpy_baseline(A, B):
    """Força bruta: Calcula tudo, inclusive os zeros."""
    start = time.time()
    C = A @ B
    return time.time() - start

def run_oxta_hybrid(A, B, active_blocks, block_size, n):
    """
    Estratégia Oxta:
    1. Ignora memória fria.
    2. Despacha blocos quentes para o motor AVX (NumPy) individualmente.
    """
    start = time.time()
    C = np.zeros((n, n), dtype=np.float32)

    # Iteramos APENAS sobre a lista de blocos ativos (Sparse Loop)
    for (bi, bk) in active_blocks:
        # Mapeamento de Coordenadas
        i_start = bi * block_size
        i_end = i_start + block_size
        k_start = bk * block_size
        k_end = k_start + block_size

        # Slicing Zero-Copy (Views)
        A_block = A[i_start:i_end, k_start:k_end]
        B_strip = B[k_start:k_end, :]  # Fatia horizontal correspondente em B

        # >>> FUSÃO: Usamos a força do NumPy apenas onde necessário <<<
        # A contribuição do bloco é somada na fatia de saída
        C[i_start:i_end, :] += A_block @ B_strip

    return time.time() - start

@njit(parallel=True, fastmath=True)
def bitwise_kernel_mask(A, B_packed, C, mask, block_size, N):
    """
    Motor Bitwise (Simulado).
    Lê pesos empacotados (4 pesos/byte) para economizar banda de memória.
    Realiza decodificação bitwise 'on-the-fly' nos registradores.
    """
    n_b_rows, n_b_cols = mask.shape

    # Paralelismo por linha de blocos
    for bi in prange(n_b_rows):
        i_start = bi * block_size

        for bk in range(n_b_cols):
            if not mask[bi, bk]: continue

            k_start = bk * block_size

            # Loop Tiling
            for i_local in range(block_size):
                i_global = i_start + i_local

                # Para cada linha de A, percorremos a linha de B (que é comprimida)
                for k_local in range(block_size):
                    k_global = k_start + k_local
                    val = A[i_global, k_global]

                    if val == 0: continue

                    # Acessar B comprimido
                    # B_packed[k, :] tem N/4 bytes.
                    # Iteramos sobre os bytes (pacotes de 4 colunas)

                    for j_pack in range(N // 4):
                        # Load 1 byte containing 4 weights
                        byte_val = B_packed[k_global, j_pack]

                        if byte_val == 0: continue # Otimização: Se byte é 0, todos pesos são 0?
                        # No nosso encoding: 00 -> 0. 11 -> 0 (padding).
                        # Se byte == 0, então w0=w1=w2=w3=0. Skip!
                        # Essa checagem economiza 4 iterações.

                        # Decodificação (Unrolling manual para velocidade)
                        # Peso 0
                        c0 = byte_val & 3
                        j0 = j_pack * 4
                        if c0 == 1:
                            C[i_global, j0] += val
                        elif c0 == 2:
                            C[i_global, j0] -= val

                        # Peso 1
                        c1 = (byte_val >> 2) & 3
                        j1 = j0 + 1
                        if c1 == 1:
                            C[i_global, j1] += val
                        elif c1 == 2:
                            C[i_global, j1] -= val

                        # Peso 2
                        c2 = (byte_val >> 4) & 3
                        j2 = j0 + 2
                        if c2 == 1:
                            C[i_global, j2] += val
                        elif c2 == 2:
                            C[i_global, j2] -= val

                        # Peso 3
                        c3 = (byte_val >> 6) & 3
                        j3 = j0 + 3
                        if c3 == 1:
                            C[i_global, j3] += val
                        elif c3 == 2:
                            C[i_global, j3] -= val

def run_oxta_bitwise(A, B_packed, active_blocks, block_size, n):
    """
    Wrapper para o motor Bitwise.
    """
    start = time.time()

    n_b = n // block_size
    mask = np.zeros((n_b, n_b), dtype=np.bool_)
    if len(active_blocks) > 0:
        mask[active_blocks[:, 0], active_blocks[:, 1]] = True

    C = np.zeros((n, n), dtype=np.float32)
    bitwise_kernel_mask(A, B_packed, C, mask, block_size, n)
    return time.time() - start

@njit(parallel=True, fastmath=True)
def structured_kernel_2_4(A, B, C, mask, block_size, N):
    """
    Kernel otimizado para Esparsidade Estruturada 2:4.
    Processa grupos de 4 elementos na dimensão K (redução).
    Simula o uso de instruções vetoriais que saltam zeros conhecidos.
    """
    n_b_rows, n_b_cols = mask.shape

    for bi in prange(n_b_rows):
        i_start = bi * block_size

        for bk in range(n_b_cols):
            if not mask[bi, bk]: continue

            k_start = bk * block_size

            # Loop Tiling
            for i_local in range(block_size):
                i_global = i_start + i_local

                # Unroll por 4 na dimensão K
                # K deve ser múltiplo de 4 (garantido pelo gerador 2:4)
                for k_local in range(0, block_size, 4):
                    k_global = k_start + k_local

                    # Carregamos 4 valores de A
                    # (Na prática hardware, carregaríamos vetor)
                    a0 = A[i_global, k_global]
                    a1 = A[i_global, k_global+1]
                    a2 = A[i_global, k_global+2]
                    a3 = A[i_global, k_global+3]

                    # Se todos A forem zero, skip (intersecção de esparsidade)
                    if a0 == 0 and a1 == 0 and a2 == 0 and a3 == 0:
                        continue

                    # Vetorização na dimensão J (Output)
                    # Para cada coluna J, processamos o grupo de 4 K's
                    for j in range(N):
                        # Carrega B (2:4 sparse)
                        # Sabemos que apenas 2 destes são não-zero.
                        # Instrução simulada: DotProductSparse

                        b0 = B[k_global, j]
                        b1 = B[k_global+1, j]
                        b2 = B[k_global+2, j]
                        b3 = B[k_global+3, j]

                        acc = 0.0
                        if b0 != 0: acc += a0 * b0
                        if b1 != 0: acc += a1 * b1
                        if b2 != 0: acc += a2 * b2
                        if b3 != 0: acc += a3 * b3

                        if acc != 0:
                            C[i_global, j] += acc

def run_oxta_structured(A, B, active_blocks, block_size, n):
    """
    Wrapper para o motor Structured 2:4.
    """
    start = time.time()
    n_b = n // block_size
    mask = np.zeros((n_b, n_b), dtype=np.bool_)
    if len(active_blocks) > 0:
        mask[active_blocks[:, 0], active_blocks[:, 1]] = True

    C = np.zeros((n, n), dtype=np.float32)
    structured_kernel_2_4(A, B, C, mask, block_size, n)
    return time.time() - start
