"""AWQ Triton Implementation.

This module contains the AWQ (Activation-Weight Quantization) implementation using Triton.
Cloned from vllm main branch (commit:cb080f32) and modified to fit roc.
Original file path: vllm/model_executor/layers/quantization/awq_triton.py
"""

# SPDX-License-Identifier: Apache-2.0

import torch
import pytest
import os
import json
import numpy as np

import itertools
import argparse
import ctypes
from typing import Optional
import aiter
from aiter import logger
from aiter import pertoken_quant, get_hip_quant
from aiter import ActivationType, QuantType, dtypes
from aiter.ops.awq_gemm_asm import *
from aiter.test_common import checkAllclose, perftest
import torch.nn.functional as F

import pandas as pd


def reverse_awq_order(tensor: torch.Tensor) -> torch.Tensor:
    """Reverse the AWQ order of the given tensor.
    
    Args:
        tensor: Input tensor to reorder
        
    Returns:
        Reordered tensor with bits masked to 4 bits
    """
    bits = 4
    AWQ_REVERSE_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
    reverse_order_tensor = torch.arange(
        tensor.shape[-1],
        dtype=torch.int32,
        device=tensor.device,
    )
    reverse_order_tensor = reverse_order_tensor.view(-1, 32 // bits)
    reverse_order_tensor = reverse_order_tensor[:, AWQ_REVERSE_ORDER]
    reverse_order_tensor = reverse_order_tensor.view(-1)

    tensor = tensor[:, reverse_order_tensor] & 0xF
    return tensor


def awq_dequantize_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
    """
    #torch.set_printoptions(profile="full")

    if group_size == -1:
        group_size = qweight.shape[0]

    bits = 4
    shifts = torch.arange(0, 32, bits, device=qzeros.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)

    scales = scales.repeat_interleave(group_size, dim=0)
    zeros = zeros.repeat_interleave(group_size, dim=0)

    #return (iweights - zeros) * scales
    print(f"[awq_dequantize_torch] iweights: shape={tuple(iweights.shape)}, dtype={iweights.dtype}, "
          f"first={(iweights.reshape(-1)[0].item() if iweights.numel() else 'NA')}")
    return (iweights - zeros) * scales, iweights, zeros, scales


def 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)

    iscales = scales.repeat_interleave(group_size, dim=0)
    zeros = zeros.repeat_interleave(group_size, dim=0)

    return (iweights - zeros) * iscales, iweights, zeros, iscales


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 print_hex(data, startY, endY, startX, endX, log, p_dtype):
    int_view = data[startY:endY, startX:endX].view(p_dtype)
    for y in range(endY-startY):
        hex_str = [hex(x.item()) for x in int_view[y]]
        print("%s: %s"%(log, hex_str))

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


@perftest(num_warmup=1, num_iters=10)  #测试次数
def asm_awq_gemm(out_asm, iweights_pack, input, zeros_pack, scales_pack):
    aiter.awq_gemm_asm(out_asm, iweights_pack, input, zeros_pack, scales_pack)
    return out_asm

# input   - [M, K]
# qweight - [K, N // 8]
# qzeros  - [K // G, N // 8]
# scales  - [K // G, N]
def test_gemm(N, K, M, G,
          dtype=dtypes.fp16,
          awq_gemm_layout="TN",
          fused_awq_gemm=True):
    print(f"[{awq_gemm_layout}] {'fused!' if fused_awq_gemm else '...'}")
    G = G if G != -1 else K
    device = "cuda"
    input_rows = M
    input_cols = K
    input_dtype = dtype
    qweight_rows = input_cols
    qweight_cols = N // 2
    scales_rows = qweight_rows // G
    scales_cols = N
    scales_dtype = dtype
    qzeros_rows = scales_rows
    qzeros_cols = qweight_cols

    torch.manual_seed(0)

    input = torch.rand((input_rows, input_cols),
                      dtype=input_dtype,
                      device=device)
    qweight = torch.randint(0,
                          torch.iinfo(torch.int8).max,
                          (qweight_rows, qweight_cols),
                          dtype=torch.int8,
                          device=device)
    qzeros = torch.randint(0,
                         torch.iinfo(torch.int8).max,
                         (qzeros_rows, qzeros_cols),
                         dtype=torch.int8,
                         device=device)
    scales = torch.rand((scales_rows, scales_cols),
                       dtype=scales_dtype,
                       device=device)

    out_asm = torch.rand((M, N),
                      dtype=input_dtype,
                      device=device)

    # output_triton = awq_gemm_triton(input, qweight, scales, qzeros)
    # assert (not torch.any(torch.isinf(output_triton))
    #         and not torch.any(torch.isnan(output_triton)))

    #torch.set_printoptions(profile="full")
    #torch.set_printoptions(profile="default")

    

    if awq_gemm_layout == "TN":
        dequantized_weights, iweights, zeros, scales = awq_dequantize_torch(qweight, scales, qzeros, G)
        dequantized_weights = dequantized_weights.transpose(0, 1).contiguous()
        iweights = iweights.transpose(0, 1).contiguous()
        zeros = zeros.transpose(0, 1).contiguous()
        scales = scales.transpose(0, 1).contiguous()
        #output_torch = F.linear(dequantized_weights, input, None)
        output_torch = torch.matmul(input, dequantized_weights.T)
        output_torch = output_torch.transpose(0, 1).contiguous()
    elif awq_gemm_layout == "NN":
        dequantized_weights, iweights, zeros, iscales = awq_post_dequant_torch(qweight, scales, qzeros, G)
        output_torch = torch.matmul(input, dequantized_weights)

    ave_t = 0.0
    if fused_awq_gemm:
        #  zeros = pack_int4_to_int8(zeros)
        #  iweights = pack_int4_to_int8(iweights)

        #  #int_view = zeros[0:32, 0:64].view(torch.uint8)
        #  #for y in range(32):
        #  #  hex_str = [hex(x.item()) for x in zeros]

        #  print(f"[L204] iweights: shape={tuple(iweights.shape)}, dtype={iweights.dtype}, "
              #  f"first={(iweights.reshape(-1)[0].item() if iweights.numel() else 'NA')}")
        #  print(f"[L204] input   : shape={tuple(input.shape)}, dtype={input.dtype}, "
              #  f"first={(input.reshape(-1)[0].item() if input.numel() else 'NA')}")
        #  aiter.awq_gemm_asm(out_asm, iweights, input, zeros, scales)

        if awq_gemm_layout == "NN":
            # out_asm = out_asm.transpose(0, 1).contiguous()
            (asm_output), ave_t = asm_awq_gemm(out_asm, qweight, input, qzeros, scales)
        else:
            iweights_pack = torch.zeros((N, K//2), dtype=torch.uint8, device=device)
            for y in range(iweights.size(0)):
                iweights_pack[y] = pack_int4_to_int8(iweights[y])

            zeros_pack = torch.zeros((N, K//2//64), dtype=torch.uint8, device=device)
            for y in range(zeros.size(0)):
                zeros_pack[y] = pack_int4_to_int8_64K(zeros[y])

            scales_pack = torch.zeros((N, K//64), dtype=scales_dtype, device=device)
            for y in range(scales.size(0)):
                for x in range(scales_pack.size(1)):
                    scales_pack[y, x] = scales[y, x*64]

            (asm_output), ave_t = asm_awq_gemm(out_asm, input, iweights_pack, zeros_pack, scales_pack)
            out_asm = out_asm.reshape(out_asm.shape[1], -1)
        # out_asm = out_asm.reshape(out_asm.shape[1], -1)
    else:
        aiter.awq_gemm_asm(out_asm, input, dequantized_weights)


    #  int_view = output_torch[0:4, 0:64].view(torch.uint16)
    #  for y in range(4):
        #  hex_str = [hex(x.item()) for x in int_view[y]]

    #  int_view = out_asm[0:4, 0:64].view(torch.uint16)
    #  for y in range(4):
        #  hex_str = [hex(x.item()) for x in int_view[y]]

    checkAllclose(output_torch, out_asm, msg="msg")

    assert (not torch.any(torch.isinf(output_torch))
            and not torch.any(torch.isnan(output_torch)))
    
    return ave_t
    # print(output_torch)
    
    # Move tensors to CPU for comparison
    # output_triton_cpu = output_triton.cpu()
    # output_torch_cpu = output_torch.cpu()

    # # Calculate the tolerance bound based on torch.testing.assert_close formula
    # atol = 1e-1
    # rtol = 1e-1
    # EXPECTED = output_torch_cpu.to(torch.float16)
    # ACTUAL = output_triton_cpu.to(torch.float16)
    # tolerance = atol + rtol * torch.abs(EXPECTED)
    # abs_diff = torch.abs(ACTUAL - EXPECTED)
    
    # # Find elements where absolute difference exceeds the tolerance
    # mask = abs_diff > tolerance
    
    '''
    if torch.any(mask):
        print("\nElements that exceed torch.testing.assert_close tolerance:")
        mismatched_indices = torch.nonzero(mask)
        for idx in mismatched_indices:
            i, j = idx[0].item(), idx[1].item()
            actual = ACTUAL[i,j]
            expected = EXPECTED[i,j]
            abs_difference = abs_diff[i,j]
            tolerance_at_point = tolerance[i,j]
            print(f"Position [{i},{j}]:")
            print(f"  Actual (Triton): {actual:.6f}")
            print(f"  Expected (Torch): {expected:.6f}")
            print(f"  |actual - expected|: {abs_difference:.6f}")
            print(f"  tolerance (atol + rtol*|expected|): {tolerance_at_point:.6f}")
    '''
            
    # Original assertion
    # torch.testing.assert_close(ACTUAL,
    #                          EXPECTED,
    #                          atol=atol,
    #                          rtol=rtol)

def test_awq_gemm_perf():
    M = [1, 16, 64]
    N = [576, 2048, 7168]
    K = [256, 3072]
    DTYPE = [torch.float16, torch.bfloat16]
    GROUP_SIZE = [64]
    LAYOUT = ["NN"]
    FUSED = ["true"]

    for m,n,k,dtype,group_size,layout,fused in  itertools.product(
        M, N, K, DTYPE, GROUP_SIZE, LAYOUT, FUSED):
        test_gemm(n, k, m, group_size, dtype, layout, fused)


def test_awq_gemm_perf_csv():
    DTYPE = [torch.float16]
    LAYOUT = ["NN"]
    FUSED = ["true"]
    gtimes = {}

    tests = pd.read_csv("aiter/configs/asm_tune/tuned_awq_gemm_NN.csv").fillna("")
    for i in range(len(tests)):
        ds = tests.iloc[i] 
        time = test_gemm(ds["N"], ds["K"], ds["M"], ds["G"], DTYPE[0], LAYOUT[0], FUSED[0])
        gtimes[i] = time / 1000.0
    timetable = pd.DataFrame.from_dict(gtimes, orient="index", columns=["gtimems"])
    timetable.to_csv("aiter/testtime.csv")

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-t", "--trans", type=str, default="NN",
                        choices=["TN", "NN"],
                        help="Matrix layout for AWQ GEMM (TN or NN)")
    parser.add_argument("-f", "--fused",
                        type=lambda x: str(x).lower() in ["true", "1", "yes"],
                        default=True,
                        help="Whether to enable fused AWQ GEMM (true/false)")
    return parser.parse_args()

if __name__ == "__main__":
    #args = parse_args()
    test_awq_gemm_perf()
    #test_awq_gemm_perf_csv()
