import math
import random
import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *
from itertools import product

def gen_args(seq_q, seq_k, head_num_q, head_num_k, head_size, has_alibi, has_mask, is_causal, use_block, return_lse, dtype):
    batch = len(seq_q)
    max_seq_q = max(seq_q)
    max_seq_k = max(seq_k)
    cu_seq_len_q = [0]
    cu_seq_len_k = [0]
    for i in range(batch):
        cu_seq_len_q.append(seq_q[i] + cu_seq_len_q[-1])
        cu_seq_len_k.append(seq_k[i] + cu_seq_len_k[-1])

    softmax_scale = 1 / math.sqrt(head_size)
    alibi_slope = None if has_alibi == False else torch.zeros((head_num_q)).uniform_(0, 0.1).to(torch.float32).mlu()
    attn_bias = None if has_mask is False else torch.randn((batch, head_num_q, max_seq_q, max_seq_k), dtype=dtype).mlu()
    total_seq_q = sum(seq_q)
    total_seq_k = sum(seq_k)
    q = torch.randn(total_seq_q, head_num_q, head_size, dtype=dtype, device="mlu")
    block_tables = None
    if use_block:
        block_size = 16
        num_blocks = (int)(batch * ((max_seq_k + block_size -1 )// block_size))
        max_num_blocks_per_seq = (max_seq_k + block_size - 1) // block_size
        cache_shape = (num_blocks, head_num_k, block_size, head_size)
        block_tables = random.sample(range(0, num_blocks), batch * max_num_blocks_per_seq)
        block_tables = torch.tensor(block_tables, dtype=torch.int32).mlu().view(batch, max_num_blocks_per_seq)
        k = torch.randn(size=cache_shape, dtype=dtype).mlu()
        v = torch.randn(size=cache_shape, dtype=dtype).mlu()
    else:
        k = torch.randn(total_seq_k, head_num_k, head_size, dtype=dtype, device="mlu")
        v = torch.randn(total_seq_k, head_num_k, head_size, dtype=dtype, device="mlu")
    tmo_out = torch.empty_like(q)
    out_lse = torch.randn(batch, head_num_q, max_seq_q, dtype = torch.float, device="mlu") if return_lse else None
    return (q, k, v, tmo_out, out_lse,
            torch.tensor(cu_seq_len_q, dtype=torch.int32, device="mlu"),
            torch.tensor(cu_seq_len_k, dtype=torch.int32, device="mlu"),
            alibi_slope, attn_bias, None, None,
            block_tables, max_seq_q, max_seq_k,
            softmax_scale, is_causal,
            -1, -1, "float", return_lse)

def gen_params(seq_q, seq_k, head_num_q, head_num_k, head_size, head_size_v, has_alibi, has_mask, use_block, dtype):
    batch = len(seq_q)
    max_seq_q = max(seq_q)
    max_seq_k = max(seq_k)
    cu_seq_len_q = [0]
    cu_seq_len_k = [0]
    for i in range(batch):
        cu_seq_len_q.append(seq_q[i] + cu_seq_len_q[-1])
        cu_seq_len_k.append(seq_k[i] + cu_seq_len_k[-1])
    cu_seq_len_q = torch.tensor(cu_seq_len_q, dtype=torch.int32, device="mlu")
    cu_seq_len_k = torch.tensor(cu_seq_len_k, dtype=torch.int32, device="mlu")

    alibi_slope = None if has_alibi == False else torch.zeros((head_num_q)).uniform_(0, 0.1).to(torch.float32).mlu()
    attn_bias = None if has_mask is False else torch.randn((batch, head_num_q, max_seq_q, max_seq_k), dtype=dtype).mlu()
    total_seq_q = sum(seq_q)
    total_seq_k = sum(seq_k)
    q = torch.randn(total_seq_q, head_num_q, head_size, dtype=dtype, device="mlu")
    block_tables = None
    if use_block:
        block_size = 16
        num_blocks = (int)(batch * ((max_seq_k + block_size -1 )// block_size))
        max_num_blocks_per_seq = (max_seq_k + block_size - 1) // block_size
        block_tables = random.sample(range(0, num_blocks), batch * max_num_blocks_per_seq)
        block_tables = torch.tensor(block_tables, dtype=torch.int32).mlu().view(batch, max_num_blocks_per_seq)
        k = torch.randn(num_blocks, head_num_k, block_size, head_size, dtype=dtype).mlu()
        v = torch.randn(num_blocks, head_num_k, block_size, head_size_v, dtype=dtype).mlu()
    else:
        k = torch.randn(total_seq_k, head_num_k, head_size, dtype=dtype, device="mlu")
        v = torch.randn(total_seq_k, head_num_k, head_size_v, dtype=dtype, device="mlu")
    return q, k, v, cu_seq_len_q, cu_seq_len_k, alibi_slope, attn_bias, block_tables

class TestFlashAttnOp(BtTestCase):
    def op_impl_base(self, *args):
        q, k, v, out, cu_seq_lens_q, cu_seq_lens_kv, alibi_slope, attn_bias, max_seq_len_q, \
        max_seq_len_kv, softmax_scale, is_causal, window_size_left, window_size_right, \
        compute_dtype, return_lse, block_tables, k_cache_quant_scale, v_cache_quant_scale = args
        is_pack = cu_seq_lens_q is not None
        has_block_table = block_tables is not None
        if has_block_table:
            assert is_pack == True
        batch = len(cu_seq_lens_q) - 1 if is_pack else q.size(0)
        head_num_q = q.size(-2)
        head_num_kv = k.size(-2) if block_tables is None else k.size(-3)
        head_size = q.size(-1)
        head_size_v = v.size(-1)
        assert head_num_q >= head_num_kv and head_num_q % head_num_kv == 0
        group = head_num_q // head_num_kv
        device = q.device
        repeat_dim = -3 if has_block_table else -2
        k_bd = torch.repeat_interleave(k, group, dim=repeat_dim)
        v_bd = torch.repeat_interleave(v, group, dim=repeat_dim)
        out_list = []
        inf = 1e6
        lse = torch.zeros(batch, head_num_q, max_seq_len_q, dtype=torch.float)
        for i in range(batch):
            q_i = q[cu_seq_lens_q[i]:cu_seq_lens_q[i+1], ...] if is_pack else q[i]
            if not has_block_table:
                k_i = k_bd[cu_seq_lens_kv[i]:cu_seq_lens_kv[i+1], ...] if is_pack else k[i]
                v_i = v_bd[cu_seq_lens_kv[i]:cu_seq_lens_kv[i+1], ...] if is_pack else v[i]
            else:
                block_table = block_tables[i]
                context_len = cu_seq_lens_kv[i+1] - cu_seq_lens_kv[i]
                block_size = k.size(-2)  # (num_block, head_num, block_size, head_size)
                table_end = (context_len + block_size - 1) // block_size
                block_ids = block_table[0 : table_end]
                keys, values = k[block_ids], v[block_ids]
                keys = torch.repeat_interleave(keys, group, dim=1)  #[num_block, head_q, block_size, head_size]
                keys = keys.transpose(1, 0).contiguous().view(head_num_q, -1, head_size) #[head_q, num_block * blcke_size, head_size]
                k_i = keys.transpose(1,0) #[num_block * blcke_size, head_q, head_size]
                k_i = k_i[0:context_len, ...] #[seq_k, head_q, head_size]
                values = torch.repeat_interleave(values, group, dim=1)
                values = values.transpose(1, 0).contiguous().view(head_num_q, -1, head_size_v)
                v_i = values.transpose(1,0)
                v_i = v_i[0:context_len, ...]
            qk = torch.einsum('qhd,khd->hqk', q_i, k_i).to(torch.float) * softmax_scale
            seq_q, seq_k = q_i.size(0), k_i.size(0)
            if alibi_slope is not None:
                slope = alibi_slope.reshape(1, head_num_q, 1, 1)
                slope_bias = torch.zeros(1, head_num_q, seq_q, seq_k).to(device=device)
                if is_causal:
                    relative_pos = torch.arange(-seq_k + 1, 1, dtype=torch.float32).to(device=device)
                    slope_bias = relative_pos * slope
                else:
                    row_idx = torch.arange(seq_q, dtype=torch.long).reshape(-1, 1)
                    col_idx = torch.arange(seq_k, dtype=torch.long)
                    relative_pos = torch.abs(row_idx + seq_k - seq_q - col_idx).to(device=device)
                    slope_bias = -slope * relative_pos.to(dtype=slope.dtype)
                qk += (slope_bias.squeeze(0))
            if is_causal:
                assert seq_q <= seq_k, "seq_q <= seq_k if causal=True"
                zeros = torch.zeros(seq_q, seq_k-seq_q, dtype=torch.float, device="mlu")
                tri = torch.full((seq_q, seq_q), -inf, dtype=torch.float, device="mlu").triu(diagonal=1)
                mask = torch.cat([zeros, tri], dim=1) # (q, k-q) + (q, q) => (q, k)
                qk += mask
            if window_size_left != -1 or window_size_right != -1:
                mask_w =  torch.full((seq_q, seq_k), -inf, dtype=torch.float, device="mlu")
                for qi in range(seq_q):
                    left = max(seq_k - seq_q + qi - window_size_left, 0) if window_size_left != -1 else 0
                    right = min(max(seq_k - seq_q + qi + window_size_right + 1, 0), seq_k) if window_size_right != -1 else seq_k
                    mask_w[qi, left:right] = 0
                qk += mask_w

            if attn_bias is not None:
                qk += attn_bias[i][:, :seq_q, :seq_k]
            if return_lse:
                lse[i][:, :seq_q] = torch.logsumexp(qk, dim=-1)
            attn = torch.softmax(qk, dim=-1, dtype=torch.float).to(q.dtype)
            qkv = torch.einsum('hqk,khd->qhd', attn, v_i)
            out_list.append(qkv)
        attn_out = torch.cat(out_list, dim=0)
        if is_pack == False:
            attn_out = attn_out.view(q.size(0), q.size(1), q.size(2), head_size_v)
        if return_lse:
            attn_out = (attn_out, lse)
        return attn_out

    def test_flash_attention(self):
        seq_len_list = [((38, 64, 128), (38, 64, 128)), ((30, 40, 50), (60, 90, 120))]
        head_num_list = [(32, 32), (32, 4)]
        use_block_list = [False, True]
        head_size_list = [(64, 64), (128, 512)]
        alibi_list = [False, True]
        mask_list = [False, True]
        causal_list = [False, True]
        dtype_list = [torch.half, torch.float]
        window_size_list = [(-1, -1), (10, -1), (8, 8)]
        mlu_name = torch.mlu.get_device_name()
        if "MLU3" in mlu_name:
            print("pagedattn is not implement on mlu370, skip it")
            use_block_list = [False]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        args = product(seq_len_list, head_num_list, head_size_list, alibi_list, mask_list, causal_list, use_block_list, window_size_list, dtype_list)
        for ((seq_q, seq_k), (head_num_q, head_num_k), (head_size, head_size_v), has_alibi, has_mask, is_causal, use_block,
             (window_size_left, window_size_right), dtype) in args:
            batch = len(seq_q)
            print("batch={}, seq_lens_q={}, seq_lens_k={}, head_num_q={}, head_num_k={}, head_size={}, head_size_v= {}, has_alibi={}, \
has_mask={}, is_causal={}, use_block={}, window_size_left={}, window_size_right={}, dtype={}, testing...".format(
              batch, seq_q, seq_k, head_num_q, head_num_k, head_size, head_size_v, has_alibi, has_mask, is_causal, use_block,
              window_size_left, window_size_right, dtype), flush=True)
            max_seq_q = max(seq_q)
            max_seq_k = max(seq_k)
            params = gen_params(seq_q, seq_k, head_num_q, head_num_k, head_size, head_size_v, has_alibi, has_mask, use_block, dtype)
            q, k, v, cu_seq_len_q, cu_seq_len_k, alibi_slope, attn_bias, block_tables = params
            softmax_scale = 1 / math.sqrt(head_size)

            torch_output = self.op_impl_base(q, k, v, None, cu_seq_len_q, cu_seq_len_k,
                            alibi_slope, attn_bias, max_seq_q, max_seq_k, softmax_scale,
                            is_causal, window_size_left, window_size_right, torch.float, False,
                            block_tables if use_block else None, None, None)
            tmo_output = ops.flash_attention(q, k, v, None, cu_seq_len_q, cu_seq_len_k,
                            alibi_slope, attn_bias, max_seq_q, max_seq_k, softmax_scale,
                            is_causal, window_size_left, window_size_right, torch.float, False,
                            block_tables if use_block else None, None, None)
            self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(),
                                    0.005, use_MSE=True, use_RAE=True)

            if use_block: # test block_tables is [batch, 1]
                block_size = max_seq_k
                num_blocks = (int)(batch * ((max_seq_k + block_size -1 )// block_size)) #batch
                max_num_blocks_per_seq = (max_seq_k + block_size - 1) // block_size #1
                block_tables = random.sample(range(0, num_blocks), batch * max_num_blocks_per_seq)
                block_tables = torch.tensor(block_tables, dtype=torch.int32).mlu().view(batch, max_num_blocks_per_seq)
                k = torch.randn(num_blocks, head_num_k, block_size, head_size, dtype=dtype).mlu()
                v = torch.randn(num_blocks, head_num_k, block_size, head_size_v, dtype=dtype).mlu()
                torch_output = self.op_impl_base(q, k, v, None, cu_seq_len_q, cu_seq_len_k,
                                alibi_slope, attn_bias, max_seq_q, max_seq_k, softmax_scale,
                                is_causal, -1, -1, torch.float, False,
                                block_tables if use_block else None, None, None)
                tmo_output = ops.flash_attention(q, k, v, None, cu_seq_len_q, cu_seq_len_k,
                                alibi_slope, attn_bias, max_seq_q, max_seq_k, softmax_scale,
                                is_causal, -1, -1, torch.float, False,
                                block_tables if use_block else None, None, None)
                self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(),
                                        0.005, use_MSE=True, use_RAE=True)

    def test_flash_attention_lse(self):
        seq_len_list =[((66, 77, 88), (77, 88, 99)), ((1024, 1024), (1024, 1024))]
        head_num_q, head_num_k = 16, 8
        head_size_list = [(64, 64), (16, 128)]
        is_causal = True
        has_alibi = True
        has_mask = True
        window_size_list = [(-1, -1), (10, -1)]
        dtype_list = [torch.half, torch.float]
        use_block_list = [False, True]
        mlu_name = torch.mlu.get_device_name()
        if "MLU3" in mlu_name:
            print("pagedattn is not implement on mlu370, skip it")
            use_block_list = [False]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        args = product(seq_len_list, head_size_list, use_block_list, window_size_list, dtype_list)
        for (seq_q, seq_k), (head_size, head_size_v), use_block, (window_size_left, window_size_right), dtype in args:
            batch = len(seq_q)
            print("batch={}, seq_lens_q={}, seq_lens_k={}, head_num_q={}, head_num_k={}, head_size={}, head_size_v={}, has_alibi={}, \
 has_mask={}, is_causal={}, use_block={}, window_size_left={}, window_size_right={}, dtype={}, testing...".format(
              batch, seq_q, seq_k, head_num_q, head_num_k, head_size, head_size_v, has_alibi, has_mask, is_causal, use_block,
              window_size_left, window_size_right, dtype), flush=True)
            max_seq_q = max(seq_q)
            max_seq_k = max(seq_k)
            params = gen_params(seq_q, seq_k, head_num_q, head_num_k, head_size, head_size_v, has_alibi, has_mask, use_block, dtype)
            q, k, v, cu_seq_len_q, cu_seq_len_k, alibi_slope, attn_bias, block_tables = params
            softmax_scale = 1 / math.sqrt(head_size)

            torch_output, torch_output_lse = self.op_impl_base(q, k, v, None, cu_seq_len_q, cu_seq_len_k,
                                        alibi_slope, attn_bias, max_seq_q, max_seq_k, softmax_scale,
                                        is_causal, window_size_left, window_size_right, torch.float, True,
                                        block_tables if use_block else None, None, None)
            tmo_output, tmo_output_lse = ops.flash_attention(q, k, v, None, cu_seq_len_q, cu_seq_len_k,
                                        alibi_slope, attn_bias, max_seq_q, max_seq_k, softmax_scale,
                                        is_causal, window_size_left, window_size_right, torch.float, True,
                                        block_tables if use_block else None, None, None)
            self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(),
                                    0.004, use_MSE=True, use_RAE=True)
            self.assertTensorsEqual(torch_output_lse.cpu(), tmo_output_lse.cpu(), 0.0006, use_MSE=True, use_RAE=True)

            if use_block: #test block_table = [batch, 1]
                block_size = max_seq_k
                num_blocks = (int)(batch * ((max_seq_k + block_size -1 )// block_size)) #batch
                max_num_blocks_per_seq = (max_seq_k + block_size - 1) // block_size #1
                block_tables = random.sample(range(0, num_blocks), batch * max_num_blocks_per_seq)
                block_tables = torch.tensor(block_tables, dtype=torch.int32).mlu().view(batch, max_num_blocks_per_seq)
                k = torch.randn(num_blocks, head_num_k, block_size, head_size, dtype=dtype).mlu()
                v = torch.randn(num_blocks, head_num_k, block_size, head_size_v, dtype=dtype).mlu()
                torch_output, torch_output_lse = self.op_impl_base(q, k, v, None, cu_seq_len_q, cu_seq_len_k,
                                            alibi_slope, attn_bias, max_seq_q, max_seq_k, softmax_scale,
                                            is_causal, window_size_left, window_size_right, torch.float, True,
                                            block_tables if use_block else None, None, None)
                tmo_output, tmo_output_lse = ops.flash_attention(q, k, v, None, cu_seq_len_q, cu_seq_len_k,
                                            alibi_slope, attn_bias, max_seq_q, max_seq_k, softmax_scale,
                                            is_causal, window_size_left, window_size_right, torch.float, True,
                                            block_tables if use_block else None, None, None)
                self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(),
                                        0.004, use_MSE=True, use_RAE=True)
                self.assertTensorsEqual(torch_output_lse.cpu(), tmo_output_lse.cpu(), 0.0006, use_MSE=True, use_RAE=True)

    def test_flash_attention_inplace(self):
        seq_q, seq_k, head_num_q, head_num_k, head_size, head_size_v = (66, 77, 88), (77, 88, 99), 16, 8, 64, 64
        use_block_list = [False, True]
        softmax_scale = 1 / math.sqrt(head_size)
        is_causal = True
        has_alibi = True
        has_mask = True
        batch = len(seq_q)
        max_seq_q = max(seq_q)
        max_seq_k = max(seq_k)
        cu_seq_len_q = [0]
        cu_seq_len_k = [0]
        window_size_left, window_size_right = -1, -1
        mlu_name = torch.mlu.get_device_name()
        if "MLU3" in mlu_name:
            print("pagedattn is not implement on mlu370, skip it")
            use_block_list = [False]
        total_seq_q = sum(seq_q)
        for use_block in use_block_list:
            print("test_flash_attention_inplace: use_block: {}, testing....".format(use_block))
            params = gen_params(seq_q, seq_k, head_num_q, head_num_k, head_size, head_size_v, has_alibi, has_mask, use_block, torch.half)
            q, k, v, cu_seq_len_q, cu_seq_len_k, alibi_slope, attn_bias, block_tables = params
            torch_output = self.op_impl_base(q, k, v, None, cu_seq_len_q, cu_seq_len_k,
                                alibi_slope, attn_bias, max_seq_q, max_seq_k, softmax_scale,
                                is_causal, window_size_left, window_size_right, torch.float, False,
                                block_tables if use_block else None,  None, None)
            tmo_output = torch.empty((total_seq_q, head_num_q, head_size_v), dtype=torch.half, device="mlu")
            ops.flash_attention(q, k, v, tmo_output, cu_seq_len_q, cu_seq_len_k,
                                alibi_slope, attn_bias, max_seq_q, max_seq_k, softmax_scale,
                                is_causal, window_size_left, window_size_right, torch.float, False,
                                block_tables if use_block else None, None, None)
            self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(),
                                    0.003, use_MSE=True, use_RAE=True)

            if use_block: #test block_table = [batch, 1]
                block_size = max_seq_k
                num_blocks = (int)(batch * ((max_seq_k + block_size -1 )// block_size)) #batch
                max_num_blocks_per_seq = (max_seq_k + block_size - 1) // block_size #1
                block_tables = random.sample(range(0, num_blocks), batch * max_num_blocks_per_seq)
                block_tables = torch.tensor(block_tables, dtype=torch.int32).mlu().view(batch, max_num_blocks_per_seq)
                k = torch.randn(num_blocks, head_num_k, block_size, head_size, dtype=torch.float16).mlu()
                v = torch.randn(num_blocks, head_num_k, block_size, head_size_v, dtype=torch.float16).mlu()
                torch_output = self.op_impl_base(q, k, v, None, cu_seq_len_q, cu_seq_len_k,
                                            alibi_slope, attn_bias, max_seq_q, max_seq_k, softmax_scale,
                                            is_causal, window_size_left, window_size_right, torch.float, False,
                                            block_tables if use_block else None, None, None)
                ops.flash_attention(q, k, v, tmo_output, cu_seq_len_q, cu_seq_len_k,
                                    alibi_slope, attn_bias, max_seq_q, max_seq_k, softmax_scale,
                                    is_causal, window_size_left, window_size_right, torch.float, False,
                                    block_tables if use_block else None, None, None)
                self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(),
                                        0.003, use_MSE=True, use_RAE=True)

    def test_inductor(self):
        seq_q, seq_k, head_num_q, head_num_k, head_size, dtype = (66, 77, 88), (77, 88, 99), 16, 8, 64, torch.float16
        use_block_list = [False, True]
        mlu_name = torch.mlu.get_device_name()
        if "MLU3" in mlu_name:
            print("pagedattn is not implement on mlu370, skip it")
            use_block_list = [False]
        alibi_list = [False, True]
        mask_list = [False, True]
        causal_list = [False, True]
        return_lse_list = [False, True]
        test_flags = product(use_block_list, alibi_list, mask_list, causal_list, return_lse_list)
        for use_block, has_alibi, has_mask, is_causal, return_lse in test_flags:
            print(f"==== use_block: {use_block}, has_alibi: {has_alibi}, has_mask: {has_mask}, is_causal: {is_causal} return_lse: {return_lse}====")
            args = gen_args(seq_q, seq_k, head_num_q, head_num_k, head_size, has_alibi, has_mask, is_causal, use_block, return_lse, dtype)
            self.base_opcheck(torch.ops.torch_mlu_ops.flash_attention, args)

if __name__ == '__main__':
    exit(run_unittest(TestFlashAttnOp))
