import numpy as np
import torch
import torch.nn as nn
import torch_npu
import torchair

torch.manual_seed(42)

import logging

from torch_npu.testing.testcase import TestCase, run_tests
from torchair import logger

logger.setLevel(logging.DEBUG)
torch._logging.set_logs(output_code=True, graph_code=True)
from vllm_ascend.utils import enable_custom_op

enable_custom_op()

import ucm_custom_ops

torch.set_printoptions(profile="full")
# torch.set_printoptions(
#     precision=4,    # 小数位数
#     threshold=10000, # 元素数量阈值（设为极大值以显示全部）
#     edgeitems=10000, # 每维显示边缘项数量
#     linewidth=200    # 每行字符数（避免换行）
# )


class TestCustomHammingDistTopK(TestCase):

    def test_hamming_dist_top_k_mla_eager(self):
        print(f"test_hamming_dist_top_k_mla_eager begin")
        print("=======================data=======================")
        device = "cpu"
        batch_size = 5
        num_head = 16  # 16 ok
        num_kv_head = 1
        head_dim_nope = 512
        head_dim_rope = 64
        head_dim = (
            head_dim_nope + head_dim_rope
        )  # 128, 256, 512, 576, 640 ok, 64, 704, 768
        compress_rate = 8
        compressed_dim = head_dim // compress_rate
        compressed_dim_nope = head_dim_nope // compress_rate
        compressed_dim_rope = head_dim_rope // compress_rate
        seqlen_q = 1
        sparse_ratio = 0.2
        chunk_size_value = 128

        seqlen_list = [8000] * batch_size
        seqlen = torch.tensor(seqlen_list, dtype=torch.int32, device=device)

        max_seq_len = max(seqlen_list)

        chunk_size_list = [chunk_size_value] * batch_size
        chunk_size = torch.tensor(chunk_size_list, dtype=torch.int32, device=device)

        top_k_list = [seq * sparse_ratio // chunk_size_list[0] for seq in seqlen_list]
        top_k = torch.tensor(top_k_list, dtype=torch.int32, device=device)

        num_chunks = seqlen // chunk_size

        block_size = 128
        num_blocks_per_seq = (seqlen + block_size - 1) // block_size  # 天花板除
        num_blocks = num_blocks_per_seq.sum().item() + 5

        qhash = torch.randint(
            255,
            (batch_size, num_head, seqlen_q, compressed_dim),
            dtype=torch.uint8,
            device=device,
        )
        qhash_zeros = torch.zeros(
            (batch_size, num_head, seqlen_q, compressed_dim_rope),
            dtype=torch.uint8,
            device=device,
        )
        qhash_pad = torch.cat((qhash, qhash_zeros), dim=-1).contiguous()

        khash_nope = torch.randint(
            255,
            (num_blocks, num_kv_head, block_size, compressed_dim_nope),
            dtype=torch.uint8,
            device=device,
        )
        khash_rope = torch.randint(
            255,
            (num_blocks, num_kv_head, block_size, compressed_dim_rope),
            dtype=torch.uint8,
            device=device,
        )
        khash = torch.cat((khash_nope, khash_rope), dim=-1).contiguous()
        khash_zeros = torch.zeros(
            (num_blocks, num_kv_head, block_size, compressed_dim_rope),
            dtype=torch.uint8,
            device=device,
        )
        khash_rope_pad = torch.cat((khash_rope, khash_zeros), dim=-1).contiguous()

        sink = 1
        recent = 4

        print(f"seqlen: {seqlen}")
        print(f"top_k: {top_k}")
        print(f"chunk_size: {chunk_size}")
        print(f"num_chunks: {num_chunks}")
        print(f"block_size: {block_size}")
        print(f"num_blocks_per_seq: {num_blocks_per_seq}")
        print(f"num_blocks: {num_blocks}")

        print(f"qhash shape: {qhash.shape}")  # torch.Size([2, 16, 1, 72])
        print(f"qhash_pad shape: {qhash_pad.shape}")  # torch.Size([1, 16, 1, 80])
        print(f"qhash_pad: {qhash_pad[0,0,0,:]}")

        print(f"khash_nope shape: {khash_nope.shape}")  # torch.Size([245, 1, 128, 64])
        print(f"khash_rope shape: {khash_rope.shape}")  # torch.Size([245, 1, 128, 8])
        print(f"khash shape: {khash.shape}")  # torch.Size([245, 1, 128, 72])
        print(
            f"khash_rope_pad shape: {khash_rope_pad.shape}"
        )  # torch.Size([245, 1, 128, 16])
        print(f"khash_rope_pad: {khash_rope_pad[0,0,0,:]}")

        print(f"max_seq_len: {max_seq_len}")
        print(f"sink: {sink}")
        print(f"recent: {recent}")

        # 初始化block_table
        max_num_blocks_per_seq = (max(seqlen_list) + block_size - 1) // block_size + 5
        block_table = torch.full(
            (len(num_blocks_per_seq), max_num_blocks_per_seq),
            fill_value=0,
            dtype=torch.int32,
        )

        shuffle = False
        start = 1  # 1
        for i, n in enumerate(num_blocks_per_seq):
            if shuffle:
                ids = torch.arange(start, start + n, dtype=torch.int32)
                idx = torch.randperm(n)
                block_table[i, :n] = ids[idx]
            else:
                block_table[i, :n] = torch.arange(start, start + n, dtype=torch.int32)
            start += n
        block_table = block_table.to(device=device)
        print(f"block_table.shape: {block_table.shape}")
        # print(f'block_table: {block_table}')
        indices = torch.zeros([batch_size, num_kv_head, 256], dtype=torch.int32)

        print(f"=======================op_hamming=======================")
        device_id = 0
        torch.npu.set_device(device_id)
        npu = f"npu:{device_id}"
        mask = torch.tensor([True, True, False, False, False])
        output_op = torch.ops._C_ucm.npu_hamming_dist_top_k(
            qhash.to(npu),
            khash_nope.to(npu),
            khash_rope_pad.to(npu),
            top_k.to(npu),
            seqlen.to(npu),
            chunk_size.to(npu),
            max_seq_len,
            sink,
            recent,
            None,
            block_table.to(npu),
            mask.to(npu),
            indices.to(npu),
        )
        # print(f'output_op shape: {output_op.shape}')  # torch.Size([2, 1, 512])
        print(f"output_op: {output_op}")  # torch.int32
        print(f"test_hamming_dist_top_k_mla_eager end")

    def test_hamming_dist_top_k_mla_graph(self):
        print(f"test_hamming_dist_top_k_mla_graph begin")
        print("=======================data=======================")
        device = "cpu"
        batch_size = 5
        num_head = 16  # 16 ok
        num_kv_head = 1
        head_dim_nope = 512
        head_dim_rope = 64
        head_dim = (
            head_dim_nope + head_dim_rope
        )  # 128, 256, 512, 576, 640 ok, 64, 704, 768
        compress_rate = 8
        compressed_dim = head_dim // compress_rate
        compressed_dim_nope = head_dim_nope // compress_rate
        compressed_dim_rope = head_dim_rope // compress_rate
        seqlen_q = 1
        sparse_ratio = 0.2
        chunk_size_value = 128

        seqlen_list = [8000] * batch_size
        seqlen = torch.tensor(seqlen_list, dtype=torch.int32, device=device)

        max_seq_len = max(seqlen_list)

        chunk_size_list = [chunk_size_value] * batch_size
        chunk_size = torch.tensor(chunk_size_list, dtype=torch.int32, device=device)

        top_k_list = [seq * sparse_ratio // chunk_size_list[0] for seq in seqlen_list]
        top_k = torch.tensor(top_k_list, dtype=torch.int32, device=device)

        num_chunks = seqlen // chunk_size

        block_size = 128
        num_blocks_per_seq = (seqlen + block_size - 1) // block_size  # 天花板除
        num_blocks = num_blocks_per_seq.sum().item() + 5

        qhash = torch.randint(
            255,
            (batch_size, num_head, seqlen_q, compressed_dim),
            dtype=torch.uint8,
            device=device,
        )
        qhash_zeros = torch.zeros(
            (batch_size, num_head, seqlen_q, compressed_dim_rope),
            dtype=torch.uint8,
            device=device,
        )
        qhash_pad = torch.cat((qhash, qhash_zeros), dim=-1).contiguous()

        khash_nope = torch.randint(
            255,
            (num_blocks, num_kv_head, block_size, compressed_dim_nope),
            dtype=torch.uint8,
            device=device,
        )
        khash_rope = torch.randint(
            255,
            (num_blocks, num_kv_head, block_size, compressed_dim_rope),
            dtype=torch.uint8,
            device=device,
        )
        khash = torch.cat((khash_nope, khash_rope), dim=-1).contiguous()
        khash_zeros = torch.zeros(
            (num_blocks, num_kv_head, block_size, compressed_dim_rope),
            dtype=torch.uint8,
            device=device,
        )
        khash_rope_pad = torch.cat((khash_rope, khash_zeros), dim=-1).contiguous()

        sink = 1
        recent = 4

        print(f"seqlen: {seqlen}")
        print(f"top_k: {top_k}")
        print(f"chunk_size: {chunk_size}")
        print(f"num_chunks: {num_chunks}")
        print(f"block_size: {block_size}")
        print(f"num_blocks_per_seq: {num_blocks_per_seq}")
        print(f"num_blocks: {num_blocks}")

        print(f"qhash shape: {qhash.shape}")  # torch.Size([2, 16, 1, 72])
        print(f"qhash_pad shape: {qhash_pad.shape}")  # torch.Size([1, 16, 1, 80])
        print(f"qhash_pad: {qhash_pad[0,0,0,:]}")

        print(f"khash_nope shape: {khash_nope.shape}")  # torch.Size([245, 1, 128, 64])
        print(f"khash_rope shape: {khash_rope.shape}")  # torch.Size([245, 1, 128, 8])
        print(f"khash shape: {khash.shape}")  # torch.Size([245, 1, 128, 72])
        print(
            f"khash_rope_pad shape: {khash_rope_pad.shape}"
        )  # torch.Size([245, 1, 128, 16])
        print(f"khash_rope_pad: {khash_rope_pad[0,0,0,:]}")

        print(f"max_seq_len: {max_seq_len}")
        print(f"sink: {sink}")
        print(f"recent: {recent}")

        # 初始化block_table
        max_num_blocks_per_seq = (max(seqlen_list) + block_size - 1) // block_size + 5
        block_table = torch.full(
            (len(num_blocks_per_seq), max_num_blocks_per_seq),
            fill_value=0,
            dtype=torch.int32,
        )

        shuffle = False
        start = 1  # 1
        for i, n in enumerate(num_blocks_per_seq):
            if shuffle:
                ids = torch.arange(start, start + n, dtype=torch.int32)
                idx = torch.randperm(n)
                block_table[i, :n] = ids[idx]
            else:
                block_table[i, :n] = torch.arange(start, start + n, dtype=torch.int32)
            start += n
        block_table = block_table.to(device=device)
        print(f"block_table.shape: {block_table.shape}")
        # print(f'block_table: {block_table}')
        indices = torch.zeros([batch_size, num_kv_head, 256], dtype=torch.int32)

        support_offload = None
        mask = torch.tensor([True, True, False, False, False])

        print(f"=======================op_hamming=======================")

        device_id = 0
        torch.npu.set_device(device_id)
        npu = f"npu:{device_id}"

        # start run custom ops
        class Network(nn.Module):
            def __init__(self):
                super(Network, self).__init__()

            def forward(
                self,
                qhash,
                khash,
                khash_rope,
                top_k,
                seqlen,
                chunk_size,
                max_seq_len,
                sink,
                recent,
                support_offload,
                block_table,
                mask,
                indices,
            ):

                out1 = torch.ops._C_ucm.npu_hamming_dist_top_k(
                    qhash,
                    khash,
                    khash_rope,
                    top_k,
                    seqlen,
                    chunk_size,
                    max_seq_len,
                    sink,
                    recent,
                    support_offload,
                    block_table,
                    mask,
                    indices,
                )

                return out1

        print(f"======================== PTA graph BEGIN ========================")
        npu_mode = Network().to("npu:%s" % device_id)
        from torchair.configs.compiler_config import CompilerConfig

        config = CompilerConfig()
        config.mode = "reduce-overhead"
        npu_backend = torchair.get_npu_backend(compiler_config=config)

        npu_mode = torch.compile(npu_mode, backend=npu_backend, dynamic=False)
        npu_out = npu_mode(
            qhash.to(npu),
            khash_nope.to(npu),
            khash_rope_pad.to(npu),
            top_k.to(npu),
            seqlen.to(npu),
            chunk_size.to(npu),
            max_seq_len,
            sink,
            recent,
            support_offload,
            block_table.to(npu),
            mask.to(npu),
            indices.to(npu),
        )
        print(f"npu_out = {npu_out}")
        print(f"test_hamming_dist_top_k_mla_graph end")


if __name__ == "__main__":
    run_tests()
