import torch
import torch.nn.functional as F
import ctypes
from typing import Optional
import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.ops.awq_gemm_asm import *
from aiter.ops.shuffle import reverse_awq_order
from aiter.ops.awq_gemm_asm import awq_gemm_asm
from aiter.ops.awq_dq_asm import awq_dq_asm
def pack_int4_to_int8(low_4bits):

    if len(low_4bits) % 2 != 0:
        low_4bits = torch.cat([low_4bits, torch.tensor([0], dtype=torch.uint8)])

    # 3. 将相邻两个低4位拼成一个 int8 值
    # 偶数索引：左移4位作为高4位；奇数索引：低4位
    packed = (low_4bits[::2]) | (low_4bits[1::2] << 4)
    packed = packed.to(torch.int8)  # 转回 int8（有符号）

    return packed
def pack_int4_to_int8_64K(low_4bits):

    if len(low_4bits) % 2 != 0:
        low_4bits = torch.cat([low_4bits, torch.tensor([0], dtype=torch.uint8)])

    # 3. 将相邻两个低4位拼成一个 int8 值
    # 偶数索引：左移4位作为高4位；奇数索引：低4位
    packed = (low_4bits[::128]) | (low_4bits[64::128] << 4)
    packed = packed.to(torch.int8)  # 转回 int8（有符号）
    
    return packed


# qweight - [K, N // 8]
# qzeros  - [K // G, N // 8]
# scales  - [K // G, N]
def asm_awq_reorder_and_repack(
    qweight: torch.Tensor,
    qzeros: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    #WARNING: Only support awq group_size=64
    N = qweight.shape[1] * 8
    K = qweight.shape[0]
    G = K // qzeros.shape[0]
    assert K // qzeros.shape[0] == 64, "[ERROR] ASM_AWQ_GEMM not support K Groupsize other than 64!"
    # assert (N % 512==0 or N==576), "[ERROR]ASM_AWQ_GEMM Not support Weight N other than 576 or multiplies of 512!"
    device = qzeros.device
    bits = 4
    shifts = torch.arange(0, 32, bits, device=device)

    iweights = torch.bitwise_right_shift(
        qweight[:, :, None],
        shifts[None, None, :],
    ).to(torch.int8)
    iweights = iweights.view(iweights.shape[0], -1)

    zeros = torch.bitwise_right_shift(
        qzeros[:, :, None],
        shifts[None, None, :],
    ).to(torch.int8)
    zeros = zeros.view(qzeros.shape[0], -1)

    zeros = reverse_awq_order(zeros)
    iweights = reverse_awq_order(iweights)

    iweights = torch.bitwise_and(iweights, (2**bits) - 1)
    zeros = torch.bitwise_and(zeros, (2**bits) - 1)

    iweights_packed = iweights.view(K, -1, 2)
    zeros_packed = zeros.view(K//G, -1, 2)

    # Repack weight to int32 and pack along the K direction
    # [K, N] -> [N, K]
    # iweights = iweights.transpose(1, 0).contiguous()
    packed_weights = torch.zeros([K, N//2], dtype=torch.int8, device=qweight.device)
    packed_zeros = torch.zeros([K//G, N//2], dtype=torch.int8, device=zeros.device)
    for i in range(2):
        packed_weights |= (iweights_packed[:, :, i].to(torch.int8) << (i * bits))
        packed_zeros |= (zeros_packed[:, :, i].to(torch.int8) << (i * bits))
    
    return packed_weights,packed_zeros

def asm_awq_post_dequant_torch(
    qweight: torch.Tensor,
    scales: torch.Tensor,
    qzeros: torch.Tensor,
    group_size: int,
) -> torch.Tensor:
    """Dequantize weights using PyTorch implementation.
    
    Args:
        qweight: Quantized weight tensor
        scales: Scale factors tensor
        qzeros: Zero points tensor
        group_size: Size of groups for quantization
        
    Returns:
        Dequantized tensor
    """
    if group_size == -1:
        group_size = qweight.shape[0]

    bits = 4
    shifts = torch.arange(0, 8, bits, device=qzeros.device)
    #只需要8 bit 展开
    iweights = torch.bitwise_right_shift(
        qweight[:, :, None],
        shifts[None, None, :],
    ).to(torch.int8)
    iweights = iweights.view(iweights.shape[0], -1)
    zeros = torch.bitwise_right_shift(
        qzeros[:, :, None],
        shifts[None, None, :],
    ).to(torch.int8)
    zeros = zeros.view(qzeros.shape[0], -1)

    iweights = torch.bitwise_and(iweights, (2**bits) - 1)
    zeros = torch.bitwise_and(zeros, (2**bits) - 1)

    scales = scales.repeat_interleave(group_size, dim=0)
    zeros = zeros.repeat_interleave(group_size, dim=0)
    return (iweights - zeros) * scales

def asm_awq_post_dequant(
    qweight: torch.Tensor,
    scales: torch.Tensor,
    qzeros: torch.Tensor,
    group_size: int,
) -> torch.Tensor:
    K = scales.shape[0] * group_size
    N = scales.shape[-1]
    # device = scales.device
    out = torch.empty((K, N), dtype=scales.dtype, device=qweight.device)
    awq_dq_asm(out, qweight, qzeros, scales)
    return out

# The inference function
# input   - [m, k]
# qweight - [n, k // 2]
# qzeros  - [k//g, n//2]
# scales  - [k//g, n]
def asm_awq_gemm_a16w4(input: torch.tensor,
               qweight: torch.tensor,
               scales: torch.tensor,
               qzeros: torch.tensor) -> torch.tensor:
    M,K = input.shape
    N = scales.shape[1]
    assert K % 256 == 0
    device = qzeros.device
    out_asm = torch.empty((M, N),
                      dtype=input.dtype,
                      device=device)
    awq_gemm_asm(out_asm, qweight, input, qzeros, scales)
    # out_asm = out_asm.reshape(out_asm.shape[1], -1)
    return out_asm