"""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.
"""

# SPDX-License-Identifier: Apache-2.0

import torch
import pytest
import os
import numpy as np
from aiter import dtypes
os.environ["AMDGCN_USE_BUFFER_OPS"] = "1"

from aiter.ops.triton.gemm_a16w4 import (
    awq_dequantize_triton,
    gemm_a16w4,
    reverse_awq_order,
    awq_reorder_and_repack,
    AWQ_TRITON_SUPPORTED_GROUP_SIZES
)

from aiter.awq_gemm_asm import (asm_awq_reorder_and_repack, asm_awq_gemm_a16w4, asm_awq_post_dequant)
from aiter.test_common import checkAllclose, perftest
import aiter
device = "cuda"

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
    """
    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

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)

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

@perftest(num_warmup=5, num_iters=10,testGraph=True)
def asm_api(input, qweight, scales, qzeros):
    fn = torch.compile(asm_awq_gemm_a16w4, backend="inductor", fullgraph= True)
    return fn(input, qweight, scales, qzeros)

@perftest(num_warmup=5, num_iters=10,testGraph=True)
def triton_api(input,
                qweight,
                scales,
                qzeros
                #split_k_iters: int
                ):
    return gemm_a16w4(input, qweight, scales, qzeros)

def test_gemm(M, N, K, G, dtype):
    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 // 8
    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.int32).max,
                          (qweight_rows, qweight_cols),
                          device=device)
    qzeros = torch.randint(0,
                         torch.iinfo(torch.int32).max,
                         (qzeros_rows, qzeros_cols),
                         device=device)
    scales = torch.rand((scales_rows, scales_cols),
                       dtype=scales_dtype,
                       device=device)

    qweight_repack, qzeros_repack = awq_reorder_and_repack(qweight, qzeros)
    output_triton,avg_triton = triton_api(input, qweight_repack, scales, qzeros_repack)
    assert (not torch.any(torch.isinf(output_triton))
            and not torch.any(torch.isnan(output_triton)))

    dequantized_weights = awq_dequantize_torch(qweight, scales, qzeros, G)
    output_torch = torch.matmul(input, dequantized_weights)
    assert (not torch.any(torch.isinf(output_torch))
            and not torch.any(torch.isnan(output_torch)))

    # ASM Version
    print(f"{input.dtype} {scales.dtype}")
    asm_qweight, asm_qzeros = asm_awq_reorder_and_repack(qweight, qzeros)
    quant_fn = torch.compile(asm_awq_post_dequant, backend="inductor", fullgraph= True)
    deq_out = quant_fn(asm_qweight,scales,asm_qzeros,64)
    deq_gemm = torch.matmul(input, deq_out)
    output_asm, avg_asm = asm_api(input, asm_qweight, scales, asm_qzeros)
    assert (not torch.any(torch.isinf(output_asm))
            and not torch.any(torch.isnan(output_asm)))
    
    # 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(dtype)
    TRITON_ACTUAL = output_triton_cpu.to(dtype)
    ASM_ACTUAL = output_asm.cpu().to(dtype)
    dq_out = deq_gemm.cpu().to(dtype)



    '''
    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}")
    '''
    # Convert avg_triton to scalar for formatting
    # avg_triton_scalar = avg_triton.item() if torch.is_tensor(avg_triton) else avg_triton

    # Original assertion
    msg = f"[TRITON_perf]{dtype=} {M=}, {K=}, {N=}, {G=}, triton_avg:{avg_triton:<8.2f}us"
    checkAllclose(EXPECTED, TRITON_ACTUAL, rtol, atol,msg)
    msg = f"[ASM_perf]{dtype=} {M=}, {K=}, {N=}, {G=}, asm_avg:{avg_asm:<8.2f}us"
    checkAllclose(EXPECTED, ASM_ACTUAL, rtol, atol,msg)
    msg = f"[DQ CHECK]"
    checkAllclose(EXPECTED, dq_out,rtol,atol,msg)


# test_gemm(16,256, 7168, 64)

for m in [1,32,64,128,256,16384]:
    for n in [7168]:
        for k in [256]:
            for dtype in [ torch.bfloat16]:
                test_gemm(m,n,k, 64, dtype)