import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *
from typing import Union, List, Tuple
from torch.nn.parameter import Parameter
from torch import Size
import os

dtype_list = [torch.half, torch.float]
if torch_mlu.mlu.is_bf16_supported():
    dtype_list.append(torch.bfloat16)
eps = 0.001

class TestFuseLayerNormOp(BtTestCase):
    def op_impl_base(self, *args):
        x, residual, gamma, beta, bias, eps, store_output_before_norm, quant_scale, out, dynamic_quant = args
        layernorm = torch.nn.LayerNorm(x.size(-1))
        layernorm.eps = eps
        layernorm.weight = Parameter(gamma)
        layernorm.bias = Parameter(beta)
        x = x + bias if bias is not None else x
        pro_input = x + residual if residual is not None else x
        output = layernorm(pro_input)
        if quant_scale is not None:
            output = (output * quant_scale).round().clamp(-128, 127).to(torch.int8)
        if out is None:
            if store_output_before_norm:
                return (output, pro_input)
            else:
                return output
        else:
            out.copy_(output)
            return out

    def test_layernorm(self):
        C = 128
        input_shape = (8, 8, 6, C)
        print("test layernorm...")
        for dtype in dtype_list:
            torch.manual_seed(1)
            input = torch.randn(input_shape, device="mlu", dtype=dtype)
            residual = torch.randn(input_shape, device="mlu", dtype=dtype)
            gamma = torch.randn(C, device="mlu", dtype=dtype)
            beta = torch.randn(C, device="mlu", dtype=dtype)
            tmo_out_0, tmo_out_1 = ops.fused_layer_norm(input, residual, gamma,
                        beta, None, eps, True, None, None, False)
            torch_out_0, torch_out_1 = self.op_impl_base(input, residual, gamma, beta, None, eps, True, None, None, False)
            self.assertTensorsEqual(tmo_out_0.cpu().float(), torch_out_0.cpu().float(), 0.005, use_MSE=True)
            self.assertTensorsEqual(tmo_out_1.cpu().float(), torch_out_1.cpu().float(), 0.005, use_MSE=True)
            # test output_inplace
            tmo_output = torch.empty(input_shape, device="mlu", dtype=dtype)
            torch_output = torch.empty_like(tmo_output)
            ops.fused_layer_norm(input, residual, gamma, beta, None, eps, False, None, tmo_output, False)
            self.op_impl_base(input, residual, gamma, beta, None, eps, False, None, torch_output, False)
            self.assertTensorsEqual(tmo_output.cpu().float(), torch_output.cpu().float(), 0.003, use_MSE=True)
            #test input_stride and output-continguous
            inputs_shape = (8, 8, 10, C)
            inputs = torch.randn(inputs_shape, device="mlu", dtype=dtype)
            input = inputs[:, :, 0:6, :]
            tmo_out = ops.fused_layer_norm(input, None, gamma, beta, None, eps, False, None, None, False)
            torch_out = self.op_impl_base(input, None, gamma, beta, None, eps, False, None, None, False)
            self.assertTensorsEqual(tmo_out.cpu().float(), torch_out.cpu().float(), 0.003, use_MSE=True)
            # has res
            tmo_out, tmo_res = ops.fused_layer_norm(input, residual, gamma, beta, None, eps, True, None, None, False)
            torch_out, torch_res = self.op_impl_base(input, residual, gamma, beta, None, eps, True, None, None, False)
            self.assertTensorsEqual(tmo_out.cpu().float(), torch_out.cpu().float(), 0.003, use_MSE=True)
            self.assertTensorsEqual(tmo_res.cpu().float(), torch_res.cpu().float(), 0.003, use_MSE=True)
            #res has stride
            res_shape = (8, 8, 16, C)
            res = torch.randn(res_shape, device="mlu", dtype=dtype)
            residual = res[..., 0:6, :]
            tmo_out, tmo_res = ops.fused_layer_norm(input, residual, gamma, beta, None, eps, True, None, None, False)
            torch_out, torch_res = self.op_impl_base(input, residual, gamma, beta, None, eps, True, None, None, False)
            self.assertTensorsEqual(tmo_out.cpu().float(), torch_out.cpu().float(), 0.003, use_MSE=True)
            self.assertTensorsEqual(tmo_res.cpu().float(), torch_res.cpu().float(), 0.003, use_MSE=True)
            #output has different stride
            outputs = torch.randn(8, 8, 12, C, dtype=dtype, device='mlu')
            output = outputs[..., 0:6, :]
            ops.fused_layer_norm(input, None, gamma, beta, None, eps, False, None, output, False)
            torch_out = self.op_impl_base(input, None, gamma, beta, None, eps, False, None, None, False)
            self.assertTensorsEqual(output.cpu().float(), torch_out.cpu().float(), 0.003, use_MSE=True)

    def test_squant_layernorm(self):
        C = 128
        input_shape = (8, 8, 6, C)
        print("test squant_layernorm...")
        for dtype in dtype_list:
            torch.manual_seed(1)
            input = torch.randn(input_shape, device="mlu", dtype=dtype)
            residual = torch.randn(input_shape, device="mlu", dtype=dtype)
            quant_scale = torch.randn(C, device="mlu", dtype=torch.float) * 30
            gamma = torch.randn(C, device="mlu", dtype=dtype)
            beta = torch.randn(C, device="mlu", dtype=dtype)
            tmo_out_0, tmo_out_1 = ops.fused_layer_norm(input, residual, gamma, beta, None, eps, True, quant_scale, None, False)
            torch_out_0, torch_out_1 = self.op_impl_base(input, residual, gamma, beta,  None, eps, True, quant_scale, None, False)
            self.assertTensorsEqual(tmo_out_0.cpu().float(), torch_out_0.cpu().float(), 0.006, use_MSE=True)
            self.assertTensorsEqual(tmo_out_1.cpu().float(), torch_out_1.cpu().float(), 0.006, use_MSE=True)

    def test_layernorm_stride(self):
        C = 128
        print("test layernorm stride input...")
        for dtype in dtype_list:
            gamma = torch.randn(C, device="mlu", dtype=dtype)
            beta = torch.randn(C, device="mlu", dtype=dtype)
            inputs_shape = (8, 8, 10, C)
            inputs = torch.randn(inputs_shape, device="mlu", dtype=dtype)
            input = inputs[:, :, 0:6, :]
            tmo_output = torch.empty(inputs.shape, device="mlu", dtype=dtype)
            tmo_output = tmo_output.as_strided(input.shape, input.stride())
            torch_out_0 = torch.empty_like(tmo_output)
            self.op_impl_base(input, None, gamma, beta, None, eps, False, None, torch_out_0, False)
            ops.fused_layer_norm(input, None, gamma, beta, None, eps, False, None, tmo_output, False)
            self.assertTensorsEqual(tmo_output.cpu().float(), torch_out_0.cpu().float(), 0.003, use_MSE=True)
            # test assumption get wrong stride when dim = 1
            input = inputs[:, :, 8:9, :]
            input = input.as_strided(input.shape, (10240, 1280, 0, 1))
            torch_out_0 = torch.empty_like(input)
            self.op_impl_base(input, None, gamma, beta, None, eps, False, None, torch_out_0, False)
            ops.fused_layer_norm(input, None, gamma, beta, None, eps, False, None, input, False)
            self.assertTensorsEqual(input.cpu().float(), torch_out_0.cpu().float(), 0.003, use_MSE=True)
            inputs_shape = (8, 8, 10, 2*C)
            inputs = torch.randn(inputs_shape, device="mlu", dtype=dtype)
            input1 = inputs[:, :, 0:2, 0:C]
            torch_out_0 = torch.empty_like(input1)
            self.op_impl_base(input1, None, gamma, beta, None, eps, False, None, torch_out_0, False)
            ops.fused_layer_norm(input1, None, gamma, beta, None, eps, False, None, input1, False)
            self.assertTensorsEqual(input1.cpu().float(), torch_out_0.cpu().float(), 0.003, use_MSE=True)
            input2 = inputs[:, :, :, 0:C]
            torch_out_0 = torch.empty_like(input2)
            self.op_impl_base(input2, None, gamma, beta, None, eps, False, None, torch_out_0, False)
            ops.fused_layer_norm(input2, None, gamma, beta, None, eps, False, None, input2, False)
            self.assertTensorsEqual(input2.cpu().float(), torch_out_0.cpu().float(), 0.003, use_MSE=True)
            #test 3-dim input
            inputs_shape = (8, 8, 2*C)
            inputs = torch.randn(inputs_shape, device="mlu", dtype=dtype)
            input1 = inputs[:, 0:2, 0:C]
            torch_out_0 = torch.empty_like(input1)
            self.op_impl_base(input1, None, gamma, beta, None, eps, False, None, torch_out_0, False)
            ops.fused_layer_norm(input1, None, gamma, beta, None, eps, False, None, input1, False)
            self.assertTensorsEqual(input1.cpu().float(), torch_out_0.cpu().float(), 0.003, use_MSE=True)
            input2 = inputs[0:2, 0:2, 0:C]
            torch_out_0 = torch.empty_like(input2)
            self.op_impl_base(input2, None, gamma, beta, None, eps, False, None, torch_out_0, False)
            ops.fused_layer_norm(input2, None, gamma, beta, None, eps, False, None, input2, False)
            self.assertTensorsEqual(input2.cpu().float(), torch_out_0.cpu().float(), 0.0032, use_MSE=True)

   # 防呆测试
    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_prevent(self):
        print("test_prevent....")
        func1 = ops.fused_layer_norm
        batch, seq_len, hidden_size = 5, 12, 512
        input = torch.randn(hidden_size, dtype=torch.half, device='mlu')
        residual = torch.randn(batch, seq_len, hidden_size, dtype=torch.half, device='mlu')
        quant_scale = torch.randn(hidden_size, dtype=torch.half, device='mlu')
        gamma = torch.randn(hidden_size, dtype=torch.half, device='mlu')
        beta = torch.randn(hidden_size, dtype=torch.half, device='mlu')
        self.assertException("input.dim() >= 2.",
                             func1, input, None, gamma, beta, None, eps, False, None, None)
        inputs = torch.randn(batch, seq_len, 2*hidden_size, dtype=torch.half, device='mlu')
        input = inputs[..., ::2]
        self.assertException("input last dim must be contiguous.",
                             func1, input, None, gamma, beta, None, eps, False, None, None)

        input = torch.randn(batch, seq_len, hidden_size, dtype=torch.half, device='mlu')
        self.assertException("layernorm mode need gamma and beta.",
                             func1, input, None, gamma, None, None, eps, False, None, None)
        gamma = torch.randn(2 * hidden_size, dtype=torch.half, device='mlu')
        beta = torch.randn(2 * hidden_size, dtype=torch.half, device='mlu')
        self.assertException("layernorm mode, gamma and beta size must be hidden_size.",
                             func1, input, None, gamma, beta, None, eps, False, None, None)

        gamma = torch.randn(hidden_size, dtype=torch.half, device='mlu')
        beta = torch.randn(hidden_size, dtype=torch.half, device='mlu')
        inputs = torch.randn(batch, seq_len, 12, hidden_size, dtype=torch.half, device='mlu')
        input = inputs[..., 6:9, :]
        self.assertException("quant_out is not support when input has stride.",
                             func1, input, None, gamma, beta, None, eps, False, quant_scale, None)
        inputs = torch.randn(batch, 12, seq_len, hidden_size, dtype=torch.half, device='mlu')
        input = inputs[:, 6:9, ...]
        self.assertException("check the strides of input.",
                             func1, input, None, gamma, beta, None, eps, False, None, input)
        input = inputs[..., 6:9, :]
        outputs = torch.randn(batch, 2*seq_len, 3, hidden_size, dtype=torch.half, device='mlu')
        output = outputs[:, :seq_len, ...]
        self.assertException("check the strides of output.",
                             func1, input, None, gamma, beta, None, eps, False, None, output)

    def test_inductor(self):
        print("test_inductor....")
        batch, seq_len, hidden_size = 5, 12, 512
        input = torch.randn(batch, seq_len, hidden_size, dtype=torch.half, device='mlu')
        residual = torch.randn(batch, seq_len, hidden_size, dtype=torch.half, device='mlu')
        quant_scale = torch.randn(hidden_size, dtype=torch.float, device='mlu')
        gamma = torch.randn(hidden_size, dtype=torch.half, device='mlu')
        beta = torch.randn(hidden_size, dtype=torch.half, device='mlu')
        bias = torch.randn(hidden_size, dtype=torch.half, device='mlu')
        residual_out = torch.empty_like(input)

        output = torch.zeros(input.shape, dtype=torch.int8, device='mlu')
        args = (input, output, residual, gamma, beta, bias, quant_scale, residual_out, None, "layernorm", eps, True, False)
        self.base_opcheck(torch.ops.torch_mlu_ops.fused_layernorm, args)
        args = (input, output, residual, gamma, beta, bias, quant_scale, None, None, "layernorm", eps, False, False)
        self.base_opcheck(torch.ops.torch_mlu_ops.fused_layernorm, args)
        output = torch.empty_like(input)
        args = (input, output, residual, gamma, beta, bias, None, None, None, "layernorm", eps, False, False)
        self.base_opcheck(torch.ops.torch_mlu_ops.fused_layernorm, args)


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