import numpy as np
from numba import njit, prange

@njit(fastmath=True)
def pack_ternary_weights(weights_ternary):
    """
    Empacota pesos ternários (int8: -1, 0, 1) em uint8 (4 pesos por byte).
    Esquema de Codificação (2 bits):
    0 (00) -> 0
    1 (01) -> 1
    2 (10) -> -1
    3 (11) -> 0 (Padding)

    Input: weights_ternary (N, N) dtype=int8
    Output: packed_weights (N, N // 4) dtype=uint8
    """
    rows, cols = weights_ternary.shape
    packed_cols = cols // 4
    packed = np.zeros((rows, packed_cols), dtype=np.uint8)

    for i in prange(rows):
        for j in range(packed_cols):
            base = j * 4
            # Carrega 4 pesos
            w0 = weights_ternary[i, base + 0]
            w1 = weights_ternary[i, base + 1]
            w2 = weights_ternary[i, base + 2]
            w3 = weights_ternary[i, base + 3]

            # Map values to 2-bit code
            # 0->00(0), 1->01(1), -1->10(2)

            # Helper logic
            c0 = 0
            if w0 == 1: c0 = 1
            elif w0 == -1: c0 = 2

            c1 = 0
            if w1 == 1: c1 = 1
            elif w1 == -1: c1 = 2

            c2 = 0
            if w2 == 1: c2 = 1
            elif w2 == -1: c2 = 2

            c3 = 0
            if w3 == 1: c3 = 1
            elif w3 == -1: c3 = 2

            # Pack: c3 << 6 | c2 << 4 | c1 << 2 | c0
            byte_val = (c3 << 6) | (c2 << 4) | (c1 << 2) | c0
            packed[i, j] = byte_val

    return packed

@njit(fastmath=True)
def unpack_ternary_weights(packed_weights, original_shape):
    """
    Desempacota para verificação.
    """
    rows, cols = original_shape
    unpacked = np.zeros((rows, cols), dtype=np.int8)

    for i in prange(rows):
        for j in range(cols // 4):
            byte_val = packed_weights[i, j]

            # Extract 2-bit codes
            c0 = byte_val & 3
            c1 = (byte_val >> 2) & 3
            c2 = (byte_val >> 4) & 3
            c3 = (byte_val >> 6) & 3

            # Map codes back to values
            # 0->0, 1->1, 2->-1
            vals = np.zeros(4, dtype=np.int8)
            codes = [c0, c1, c2, c3]

            for k in range(4):
                c = codes[k]
                if c == 1:
                    unpacked[i, j*4 + k] = 1
                elif c == 2:
                    unpacked[i, j*4 + k] = -1
                else:
                    unpacked[i, j*4 + k] = 0

    return unpacked
