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 TestFuseRmsNormOp(BtTestCase):
    def run_gen_case(self, dic):
        dump_data = dic.pop('dump_data')
        if dump_data:
            self.launch(*dic.values())
        else:
            x = create_tensor_from_dic(dic['x'])
            residual = None if dic['residual']['data'] is None else create_tensor_from_dic(dic['residual'])
            gamma = None if dic['gamma']['data'] is None else create_tensor_from_dic(dic['gamma'])
            beta = None if dic['beta']['data'] is None else create_tensor_from_dic(dic['beta'])
            bias = None if dic['bias']['data'] is None else create_tensor_from_dic(dic['bias'])
            eps = dic['eps']['data']
            store_output_before_norm = dic['store_output_before_norm']['data']
            quant_scale = None if dic['quant_scale']['data'] is None else create_tensor_from_dic(dic['quant_scale'])
            out = None if dic['out']['data'] is None else create_tensor_from_dic(dic['out'])
            dynamic_quant = dic['dynamic_quant']['data']
            self.launch(x, residual, gamma, beta, bias, eps, store_output_before_norm, quant_scale, out, dynamic_quant)

    def launch(self, *args):
        args = list(args)
        base_out = None if args[-2] is None else torch.empty_like(args[-2])
        base_input = args[0].clone() if args[-2] is not None and args[0] is args[-2] else args[0]
        tmo_out = ops.fused_rms_norm(*args)
        args[0] = base_input
        args[-2] = base_out
        torch_out = self.op_impl_base(*args)
        if tmo_out.__class__ in (list, tuple):
            for o1, o2 in zip(tmo_out, torch_out):
                self.assertTensorsEqual(o1.cpu().float(), o2.cpu().float(), 0.007, use_MSE=True)
        else:
            self.assertTensorsEqual(tmo_out.cpu().float(), torch_out.cpu().float(), 0.007, use_MSE=True)

    def op_impl_base(self, *args):
        x, residual, gamma, beta, bias, eps, store_output_before_norm, quant_scale, out, dynamic_quant = args
        x = x + bias if bias is not None else x
        pro_input = x + residual if residual is not None else x
        store_input = pro_input
        pro_input = pro_input.to(torch.float32)
        variance = pro_input.pow(2).mean(-1, keepdim=True)
        pro_input = pro_input * torch.rsqrt(variance + eps)
        if gamma.dtype in [torch.float16, torch.bfloat16]:
            pro_input = pro_input.to(gamma.dtype)
        output = gamma * 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, store_input)
            else:
                return output
        else:
            out.copy_(output)
            return out

    def test_rmsnorm(self):
        C = 128
        input_shape = (8, 8, 6, C)
        print("test rmsnorm...")
        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)
            self.launch(input, residual, gamma, None, None, eps, True, None, None, False)
            # test inplace output
            output = torch.empty(input_shape, device="mlu", dtype=dtype)
            self.launch(input, residual, gamma, None, None, eps, False, None, output, False)
            inputs_shape = (8, 8, 10, C)
            inputs = torch.randn(inputs_shape, device="mlu", dtype=dtype)
            input = inputs[:, :, 0:6, :]
            self.launch(input, None, gamma, None, None, eps, False, None, None, False)
            self.launch(input, residual, gamma, None, None, eps, True, None, None, False)
            #res has stride
            res_shape = (8, 8, 16, C)
            res = torch.randn(res_shape, device="mlu", dtype=dtype)
            residual = res[..., 0:6, :]
            self.launch(input, residual, gamma, None, None, eps, True, None, None, False)
            #output has different stride
            outputs = torch.randn(8, 8, 12, C, dtype=dtype, device='mlu')
            output = outputs[..., 0:6, :]
            self.launch(input, residual, gamma, None, None, eps, False, None, output, False)

    def test_squant_rmsnorm(self):
        C = 128
        input_shape = (8, 8, 6, C)
        print("test squant_rmsnorm...")
        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)
            self.launch(input, residual, gamma, None, None, eps, True, quant_scale, None, False)
            # test one output
            self.launch(input, None, gamma, None, None, eps, False, quant_scale, None, False)

    def test_rmsnorm_stride(self):
        C = 128
        print("test rmsnorm stride input...")
        for dtype in dtype_list:
            gamma = 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, :]
            output = torch.empty(inputs.shape, device="mlu", dtype=dtype)
            output = output.as_strided(input.shape, input.stride())
            self.launch(input,  None, gamma, None, None, eps, False, None, output, False)
            # test assumption get wrong stride when dim = 1
            input = inputs[:, :, 0 :1, :]
            input = input.as_strided(input.shape, (10240, 1280, 0, 1))
            self.launch(input, None, gamma, None, None, eps, False, None, input, False)
            inputs_shape = (8, 8, 10, 2*C)
            inputs = torch.randn(inputs_shape, device="mlu", dtype=dtype)
            input1 = inputs[:, :, 0:2, 0:C]
            self.launch(input1, None, gamma, None, None, eps, False, None, input1, False)
            input2 = inputs[:, :, :, 0:C]
            self.launch(input2, None, gamma, None, None, eps, False, None, input2, False)
            #tset 3-dim input
            inputs_shape = (8, 8, 2*C)
            inputs = torch.randn(inputs_shape, device="mlu", dtype=dtype)
            input1 = inputs[:, 0:2, 0:C]
            self.launch(input1, None, gamma, None, None, eps, False, None, input1, False)
            input2 = inputs[0:2, 0:2, 0:C]
            self.launch(input2, None, gamma, None, None, eps, False, None, input2, False)


    # 防呆测试
    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_prevent(self):
        func1 = ops.fused_rms_norm
        batch, seq_len, hidden_size = 5, 12, 512
        input = torch.randn(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')
        self.assertException("input.dim() >= 2.",
                             func1, input, None, gamma, None, 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, None, None, eps, False, None, None)

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

        gamma = 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, None, 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, None, None, eps, False, None, input)
        input = inputs[..., 6:9, :]
        outputs = torch.randn(batch, 2 * seq_len, 12, hidden_size, dtype=torch.half, device='mlu')
        output = outputs[:, 0:seq_len, 6:9, :]
        self.assertException("check the strides of output.",
                             func1, input, None, gamma, None, None, eps, False, None, output)


    def test_inductor(self):
        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, "rmsnorm", 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, "rmsnorm", 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, "rmsnorm", eps, False, False)
        self.base_opcheck(torch.ops.torch_mlu_ops.fused_layernorm, args)

if __name__ == '__main__':
    run_unittest(TestFuseRmsNormOp)
