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

from torch_npu.testing.testcase import TestCase, run_tests
from torch_npu.testing.common_utils import create_common_tensor
from torch_npu.testing.common_utils import SkipIfNotGteCANNVersion


class TestLayerNorm(TestCase):
    @SkipIfNotGteCANNVersion("9.0.0")
    def test_c10_layer_norm(self):
        # test that we can call c10 ops and they return a reasonable result
        X = torch.rand(5, 5, dtype=torch.float, device="cpu")
        X = X.to("npu")
        weight = torch.rand(*X.size()[1:], dtype=torch.float, device="cpu")
        weight = weight.to("npu")
        bias = torch.rand(*X.size()[1:], dtype=torch.float, device="cpu")
        bias = bias.to("npu")
        epsilon = 1e-4

        expected_norm = torch.nn.functional.layer_norm(
            X, X.size()[1:], weight=weight, bias=bias, eps=epsilon)
        expected_norm_cpu = torch.nn.functional.layer_norm(
            X.cpu(), X.size()[1:], weight=weight.cpu(), bias=bias.cpu(), eps=epsilon)
        self.assertRtolEqual(expected_norm.cpu().numpy(), expected_norm_cpu.numpy())

    def cpu_op_exec(self, input1):
        m = nn.LayerNorm(input1.size()[1:])
        output = m(input1)
        return output

    def npu_op_exec(self, input1):
        m = nn.LayerNorm(input1.size()[1:]).npu()
        output = m(input1)
        output = output.to("cpu")
        return output

    @SkipIfNotGteCANNVersion("9.0.0")
    def test_layer_norm_shape_format(self):
        shape_format = [
            [np.float32, 0, (64, 10)],
            [np.float32, 0, (256, 2048, 7, 7)],
            [np.float32, 0, (32, 1, 3, 3)],
            [np.float32, 0, (10, 128)],
            [np.float32, 2, (46, 16)],
            [np.float32, 3, (2, 2, 2)],
            [np.float32, 29, (3, 4, 5, 6)]
        ]
        for item in shape_format:
            # Ascend950 does not support FORMAT_NC1HWC0 (3) and FORMAT_FRACTAL_NZ (29)
            device_name = torch_npu.npu.get_device_name()
            if "Ascend950" in device_name and item[1] in (3, 29):
                continue
            cpu_input, npu_input = create_common_tensor(item, 1, 100)
            cpu_output = self.cpu_op_exec(cpu_input)
            npu_output = self.npu_op_exec(npu_input)
            self.assertRtolEqual(cpu_output.detach().numpy(), npu_output.detach().numpy())

    @SkipIfNotGteCANNVersion("9.0.0")
    def test_layer_norm_float16_format(self):
        shape_format = [
            [np.float16, 0, (64, 10)],
            [np.float16, 0, (256, 2048, 7, 7)],
            [np.float16, 0, (32, 1, 3, 3)],
            [np.float16, 0, (10, 128)],
            [np.float16, 2, (46, 16)],
            [np.float16, 3, (2, 2, 2)],
            [np.float16, 29, (3, 4, 5, 6)]
        ]
        for item in shape_format:
            # Ascend950 does not support FORMAT_NC1HWC0 (3) and FORMAT_FRACTAL_NZ (29)
            device_name = torch_npu.npu.get_device_name()
            if "Ascend950" in device_name and item[1] in (3, 29):
                continue
            cpu_input, npu_input = create_common_tensor(item, 1, 10)
            cpu_input = cpu_input.to(torch.float32)
            cpu_output = self.cpu_op_exec(cpu_input)
            npu_output = self.npu_op_exec(npu_input)
            cpu_output = cpu_output.to(torch.float16)
            self.assertRtolEqual(cpu_output.detach().numpy(), npu_output.detach().numpy())

    @SkipIfNotGteCANNVersion("9.0.0")
    def test_layer_norm_case_in_trocr(self):
        cpu_input = torch.rand(10, 1, 1024).uniform_(-22., 66.).half()
        cpu_weight = torch.rand(1024).uniform_(0.5, 1.1).half()
        cpu_bias = torch.rand(1024).uniform_(-0.1, 0.1).half()
        npu_input = cpu_input.npu()
        npu_weight = cpu_weight.npu()
        npu_bias = cpu_bias.npu()
        normalized_shape = (1024,)
        eps = 1e-05

        cpu_out1 = torch.layer_norm(cpu_input.float(), normalized_shape, cpu_weight.float(), cpu_bias.float(),
                                    eps, torch.backends.cudnn.enabled).half()
        npu_out1 = torch.layer_norm(npu_input, normalized_shape, npu_weight, npu_bias,
                                    eps, torch.backends.cudnn.enabled)
        self.assertRtolEqual(cpu_out1, npu_out1.cpu())

        cpu_out2 = torch.layer_norm(cpu_input.float(), normalized_shape, cpu_weight.float(), cpu_bias.float(),
                                    eps, torch.backends.cudnn.enabled).half()
        npu_out2 = torch.layer_norm(npu_input, normalized_shape, npu_weight, npu_bias,
                                    eps, torch.backends.cudnn.enabled)
        self.assertRtolEqual(cpu_out2, npu_out2.cpu())

    @SkipIfNotGteCANNVersion("9.0.0")
    def test_layer_norm_abnormal_input(self):
        npu_input = torch.randn(10, 5).npu()
        with self.assertRaises(RuntimeError) as cm:
            output = nn.functional.layer_norm(npu_input, [2, 3, 3])
        exception = cm.exception
        self.assertTrue("Given normalized_shape=[2, 3, 3], expected input with shape [*, 2, 3, 3], but got input of size[10, 5]" in str(exception))


    @SkipIfNotGteCANNVersion("9.0.0")
    def test_layer_norm_compatible_impl_switch(self):
        shape_format = [
            [np.float32, 0, (64, 10)],
            [np.float32, 0, (256, 2048, 7, 7)],
            [np.float16, 0, (10, 128)],
            [np.float16, 0, (46, 16)],
        ]

        try:
            for item in shape_format:
                cpu_input, npu_input = create_common_tensor(item, 1, 100)
                normalized_shape = cpu_input.size()[1:]

                cpu_output = torch.nn.functional.layer_norm(
                    cpu_input.float(), normalized_shape)
                if item[0] == np.float16:
                    cpu_output = cpu_output.to(torch.float16)

                torch_npu.npu.use_compatible_impl(False)
                npu_out_no_compat = torch.nn.functional.layer_norm(
                    npu_input, normalized_shape)
                self.assertRtolEqual(

                    cpu_output.detach().numpy(),
                    npu_out_no_compat.cpu().detach().numpy())

                torch_npu.npu.use_compatible_impl(True)
                npu_out_with_compat = torch.nn.functional.layer_norm(
                    npu_input, normalized_shape)
                self.assertRtolEqual(
                    cpu_output.detach().numpy(),
                    npu_out_with_compat.cpu().detach().numpy())
        finally:
            torch_npu.npu.use_compatible_impl(True)


if __name__ == "__main__":
    torch_npu.npu.use_compatible_impl(True)
    run_tests()