import torch
import numpy as np


def B_global_16x32_to_shared_load_16x32_layout(i, j):
    """
         stride * 8 * (tx // HALF_WARP_expr)
                + (tx % 8) * stride
                + 16 * ((tx % HALF_WARP_expr) // 8)
    """
    thread_id = i * 2 + j // 16
    row = (thread_id // 16) * 8 + (thread_id % 8)
    col = (j % 16) + 16 * ((thread_id % 16) // 8)
    return row, col


def permutate_weight_fastest(weight):
    wmma_n = 16
    wmma_k = 32
    N = weight.shape[0]
    K = weight.shape[1]
    
    # Create a lookup table for the permutation
    mapping = np.zeros((wmma_n, wmma_k, 2), dtype=int)
    for ii in range(wmma_n):
        for jj in range(wmma_k):
            mapping[ii, jj] = B_global_16x32_to_shared_load_16x32_layout(ii, jj)
    
    # Reshape weight for the final format
    permutated_weight = np.zeros((N // wmma_n, K // wmma_k, wmma_n, wmma_k), dtype="int8")
    
    # Use advanced indexing for the entire operation
    i_indices = np.arange(N // wmma_n)[:, np.newaxis, np.newaxis, np.newaxis]
    j_indices = np.arange(K // wmma_k)[np.newaxis, :, np.newaxis, np.newaxis]
    
    # Create the source indices
    src_i = i_indices * wmma_n + mapping[:, :, 0]
    src_j = j_indices * wmma_k + mapping[:, :, 1]
    
    # Extract and reshape in one go
    permutated_weight = weight[src_i, src_j]
    
    return permutated_weight


def compress_int2_to_int8(int2_weight):
    int8_weight = np.zeros(
        (*int2_weight.shape[:-1], int2_weight.shape[-1] // 4), dtype=np.int8
    )
    for j in range(int2_weight.shape[-1] // 4):
        for k in range(4):
            int8_weight[:, :, :, j] |= int2_weight[:, :, :, j * 4 + k] << (k * 2)
    return int8_weight


def interleave_weight_int8(qweight, nbits=2):\
    # reinterpret the data type of qweight to int32
    # shift = [ 0,  8, 16, 24,  2, 10, 18, 26,  4, 12, 20, 28,  6, 14, 22, 30]
    # index: [ 0,  4,  8, 12,  1,  5,  9, 13,  2,  6, 10, 14,  3,  7, 11, 15]
    qweight = qweight.view(np.int32)
    new_qweight = np.zeros_like(qweight)
    bits_stride = 8
    mask = (1 << nbits) - 1  # for 4bit the val is 0x0000000f
    num_groups = 32 // bits_stride # 4
    elems_per_group = bits_stride // nbits  # 4
    for i in range(num_groups):
        for j in range(elems_per_group):
            offset = i * elems_per_group + j
            shift = (offset % num_groups) * bits_stride + (offset // num_groups) * nbits

            new_qweight |= ((qweight >> (nbits * offset)) & mask) << shift
    return new_qweight.view(np.int8)



def convert_weight_int8_to_int2(weight):
    N = weight.shape[0]
    K = weight.shape[1]

    weight = weight+2
    
    weight = weight.cpu().numpy()

    # print(weight)
    # print(torch.max(weight), torch.min(weight))

    # permutated_weight_slow = permutate_weight(weight)
    permutated_weight = permutate_weight_fastest(weight)
    # assert np.all(permutated_weight_slow == permutated_weight)
    # print("Permutation is correct")
    compressed_weight = compress_int2_to_int8(permutated_weight)
    interleaved_weight = interleave_weight_int8(compressed_weight, 2)

    ret = torch.from_numpy(interleaved_weight)

    ret = torch.reshape(ret, (N, K // 4))

    return ret