import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *
from torch.nn.parameter import Parameter
from torch.nn import functional as F
import torch.nn as nn

class TestFusedNormAttnProjOp(BtTestCase):
    def op_impl_base(self, *args):
        input, q_weight, q_bias, k_weight, k_bias, v_weight, v_bias, norm_weight, \
            norm_bias, eps, out_layout, head_size, norm_out = args
        input_size = input.size(-1)
        layernorm = torch.nn.LayerNorm(input_size)
        layernorm.eps = eps
        layernorm.weight = nn.Parameter(norm_weight)
        layernorm.bias = nn.Parameter(norm_bias)
        layernorm_out = layernorm(input)
        q_out = torch.matmul(layernorm_out, q_weight.permute(1, 0)) + q_bias
        if k_weight is not None:
            k_out = torch.matmul(layernorm_out, k_weight.permute(1, 0)) + k_bias
            v_out = torch.matmul(layernorm_out, v_weight.permute(1, 0)) + v_bias

        if out_layout == 'nhtc':
            batch, seq, _ = input.shape
            hidden_size_q = q_weight.size(0)
            q_head = hidden_size_q // head_size
            q_out = q_out.reshape(batch, seq, q_head, head_size).transpose(1, 2)
            if k_weight is not None:
                hidden_size_kv = k_weight.size(0)
                kv_head = hidden_size_kv // head_size
                k_out = k_out.reshape(batch, seq, kv_head, head_size).transpose(1, 2)
                v_out = v_out.reshape(batch, seq, kv_head, head_size).transpose(1, 2)
        outs = (q_out,) if k_weight is None else (q_out, k_out, v_out,)
        if norm_out is True:
            outs += (layernorm_out,)
        return outs

    def test_attn_proj(self):
        dtype_list = [torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        for dtype in dtype_list:
            N, T, input_size, hidden_size, head_size, eps, alpha, beta = 4, 16, 512, 768, 64, 1e-5, 0.5, 0.3
            print("N: {}, T: {}, input_size: {}, hidden_size: {}, testing...".format(
                    N, T, input_size, hidden_size), flush=True)
            input = torch.randn(N, T, input_size, dtype=dtype, device="mlu")
            weight = torch.randn(hidden_size * 3, input_size, dtype=dtype, device="mlu")
            bias = torch.randn(hidden_size * 3, dtype=dtype, device="mlu")
            norm_weight = torch.randn(input_size, dtype=dtype, device="mlu")
            norm_bias = torch.randn(input_size, dtype=dtype, device="mlu")
            residual = torch.randn(N, T, hidden_size, dtype=dtype, device="mlu")
            weights = torch.chunk(weight, 3)
            biass = torch.chunk(bias, 3)
            # test pre_attn_proj
            print("test pre_attn_proj...")
            out_torch = self.op_impl_base(input,
                                          weights[0], biass[0],
                                          weights[1], biass[1],
                                          weights[2], biass[2],
                                          norm_weight, norm_bias,
                                          eps, 'nthc', head_size, True)
            out_tmo = ops.fused_norm_attention_project(input,
                                                        weights[0], biass[0],
                                                        weights[1], biass[1],
                                                        weights[2], biass[2],
                                                        norm_weight, norm_bias,
                                                        eps, 'nthc', head_size, True)
            for o1, o2 in list(zip(out_torch, out_tmo)):
                self.assertTensorsEqual(o1.cpu().float(), o2.cpu().float(),
                                        0.003, use_MSE=True, use_RAE=True)

    def test_inductor(self):
        N, T, input_size, hidden_size, head_size, eps, alpha, beta, dtype = 4, 16, 512, 768, 64, 1e-5, 0.5, 0.3, torch.half
        input = torch.randn(N, T, input_size, dtype=dtype, device="mlu")
        weight = torch.randn(hidden_size * 3, input_size, dtype=dtype, device="mlu")
        bias = torch.randn(hidden_size * 3, dtype=dtype, device="mlu")
        norm_weight = torch.randn(input_size, dtype=dtype, device="mlu")
        norm_bias = torch.randn(input_size, dtype=dtype, device="mlu")
        weights = torch.chunk(weight, 3)
        biass = torch.chunk(bias, 3)
        args = (input, weights[0], biass[0], weights[1], biass[1], weights[2], biass[2],
                norm_weight, norm_bias, None, "nhtc", head_size, eps, alpha,
                beta, True)
        self.base_opcheck(torch.ops.torch_mlu_ops.attention_project, args)

        args = (input, weights[0], biass[0], weights[1], biass[1], weights[2], biass[2],
                norm_weight, norm_bias, None, "nthc", head_size, eps, alpha,
                beta, True)
        self.base_opcheck(torch.ops.torch_mlu_ops.attention_project, args)

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