import os
import sys
import unittest
import math
from unittest.mock import patch, MagicMock, Mock

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
import numpy as np

class TestAdapterNPU(unittest.TestCase):
    def setUp(self):
        self.adapter_module = None
        try:
            import adapter
            self.adapter_module = adapter
        except ImportError:
            self.skipTest("adapter module not found")
    
    @patch('adapter.HAS_NPU', True)
    @patch('adapter.torch_npu')
    @patch('adapter.torchair')
    def test_has_npu_with_npu_available(self, mock_torchair, mock_torch_npu):
        import adapter
        self.assertTrue(adapter.has_npu())
    
    @patch('adapter.HAS_NPU', False)
    def test_has_npu_without_npu(self):
        import adapter
        self.assertFalse(adapter.has_npu())
    
    @patch('adapter.HAS_NPU', True)
    @patch('builtins.print')
    def test_adapt_to_accelerator_device_npu(self, mock_print):
        import adapter
        adapter.adapt_to_accelerator_device()
        mock_print.assert_called_with("load adapter for accelerator device, now support to run on Ascend NPU")
    
    @patch('adapter.HAS_NPU', False)
    @patch('builtins.print')
    def test_adapt_to_accelerator_device_gpu(self, mock_print):
        import adapter
        adapter.adapt_to_accelerator_device()
        mock_print.assert_called_with("load adapter for accelerator device, now support to run on Nvidia GPU")
    
    @patch('adapter.HAS_NPU', True)
    @patch('adapter.dist')
    @patch.dict(os.environ, {'LOCAL_RANK': '0'})
    def test_init_distributed_adapter_npu(self, mock_dist):
        import adapter
        adapter.init_distributed_adapter()
        mock_dist.init_process_group.assert_called_with(
            backend="cpu:gloo,npu:hccl",
            init_method="env://"
        )
    
    @patch('adapter.HAS_NPU', False)
    @patch('adapter.dist')
    @patch.dict(os.environ, {'LOCAL_RANK': '1'})
    def test_init_distributed_adapter_gpu(self, mock_dist):
        import adapter
        adapter.init_distributed_adapter()
        mock_dist.init_process_group.assert_called_with(
            backend="nccl",
            device_id=torch.device("cuda", 1)
        )
    
    @patch.dict(os.environ, {}, clear=True)
    def test_init_distributed_adapter_no_local_rank(self):
        import adapter
        with patch('adapter.dist') as mock_dist:
            adapter.init_distributed_adapter()
            mock_dist.init_process_group.assert_not_called()
    
    @patch('torch.use_deterministic_algorithms')
    def test_deterministic_on(self, mock_use_deterministic):
        import adapter
        try:
            from msprobe.pytorch import seed_all
            with patch('msprobe.pytorch.seed_all') as mock_seed_all:
                adapter.deterministic_on()
                mock_use_deterministic.assert_called_once_with(True)
                mock_seed_all.assert_called_once_with(seed=42, mode=True, rm_dropout=True)
        except ImportError:
            mock_use_deterministic.assert_called_once_with(True)
    
    @patch('adapter.HAS_NPU', True)
    @patch('adapter.LongContextAttention')
    @patch('adapter.AttnType')
    def test_get_longcontext_attention_npu(self, mock_attn_type, mock_long_context):
        import adapter
        mock_attn_type.NPU = "npu_attn"
        result = adapter.get_longcontext_attention()
        mock_long_context.assert_called_with(
            ring_impl_type="basic_npu",
            attn_type="npu_attn"
        )
    
    @patch('adapter.HAS_NPU', False)
    @patch('adapter.LongContextAttention')
    @patch('adapter.AttnType')
    def test_get_longcontext_attention_gpu(self, mock_attn_type, mock_long_context):
        import adapter
        mock_attn_type.FA3 = "fa3_attn"
        result = adapter.get_longcontext_attention()
        mock_long_context.assert_called_with(
            ring_impl_type="basic",
            attn_type="fa3_attn"
        )
    
    @patch('adapter.HAS_NPU', True)
    @patch('adapter.torchair')
    def test_get_compiler_backend_npu(self, mock_torchair):
        import adapter
        mock_torchair.get_npu_backend.return_value = "npu_backend"
        result = adapter.get_compiler_backend()
        self.assertEqual(result, "npu_backend")
        mock_torchair.get_npu_backend.assert_called_once()
    
    @patch('adapter.HAS_NPU', False)
    def test_get_compiler_backend_gpu(self):
        import adapter
        result = adapter.get_compiler_backend()
        self.assertEqual(result, "inductor")
    
    @patch('adapter.HAS_NPU', True)
    def test_contiguous_for_channels_last_3d_memory_format_npu(self):
        import adapter
        tensor = torch.randn(2, 3, 4, 5, 6)
        result = adapter.contiguous_for_channels_last_3d_memory_format(tensor)
        self.assertTrue(result.is_contiguous())
    
    @patch('adapter.HAS_NPU', False)
    def test_contiguous_for_channels_last_3d_memory_format_gpu(self):
        import adapter
        tensor = torch.randn(2, 3, 4, 5, 6)
        result = adapter.contiguous_for_channels_last_3d_memory_format(tensor)
        self.assertTrue(result.is_contiguous(memory_format=torch.channels_last_3d))
    
    @patch('adapter.HAS_NPU', True)
    @patch('adapter.torch_npu')
    def test_adapter_gelu_npu(self, mock_torch_npu):
        mock_torch_npu.fast_gelu.return_value = torch.tensor([0.5, 0.8])
        import adapter
        x = torch.tensor([1.0, 2.0])
        result = adapter.adapter_gelu(x)
        mock_torch_npu.fast_gelu.assert_called_once_with(x)
    
    @patch('adapter.HAS_NPU', False)
    def test_adapter_gelu_gpu(self):
        import adapter
        x = torch.tensor([0.0, 1.0, -1.0, 2.0])
        result = adapter.adapter_gelu(x)
        expected = 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
        self.assertTrue(torch.allclose(result, expected, atol=1e-6))
    
    @patch('adapter.HAS_NPU', True)
    @patch('adapter.torch_npu')
    def test_adpater_rmsnorm_npu(self, mock_torch_npu):
        mock_torch_npu.npu_rms_norm.return_value = (torch.tensor([1.0, 2.0]), None)
        import adapter
        x = torch.randn(4, 8)
        weight = torch.randn(8)
        eps = 1e-5
        result = adapter.adpater_rmsnorm(x, weight, eps)
        mock_torch_npu.npu_rms_norm.assert_called_once_with(x, weight, epsilon=eps)
    
    @patch('adapter.HAS_NPU', False)
    def test_adpater_rmsnorm_gpu(self):
        import adapter
        x = torch.randn(4, 8)
        weight = torch.randn(8)
        eps = 1e-5
        result = adapter.adpater_rmsnorm(x, weight, eps)
        x_float = x.float()
        expected = x_float * torch.rsqrt(x_float.pow(2).mean(dim=-1, keepdim=True) + eps)
        if weight.dtype in [torch.float16, torch.bfloat16]:
            expected = expected.type_as(weight)
        expected = weight * expected
        self.assertTrue(torch.allclose(result, expected, atol=1e-6))
    
    @patch('adapter.HAS_NPU', True)
    @patch('adapter.torch_npu')
    def test_apply_rotary_pos_emb_adapter_npu(self, mock_torch_npu):
        mock_torch_npu.npu_rotary_mul.side_effect = [
            torch.randn(2, 4, 8, 16),
            torch.randn(2, 4, 8, 16)
        ]
        import adapter
        q = torch.randn(2, 4, 8, 16)
        k = torch.randn(2, 4, 8, 16)
        cos = torch.randn(2, 8, 16)
        sin = torch.randn(2, 8, 16)
        result_q, result_k = adapter.apply_rotary_pos_emb_adapter(q, k, cos, sin)
        self.assertEqual(mock_torch_npu.npu_rotary_mul.call_count, 2)
    
    @patch('adapter.HAS_NPU', False)
    def test_apply_rotary_pos_emb_adapter_gpu(self):
        import adapter
        q = torch.randn(2, 4, 8, 16)
        k = torch.randn(2, 4, 8, 16)
        cos = torch.randn(2, 8, 16)
        sin = torch.randn(2, 8, 16)
        result_q, result_k = adapter.apply_rotary_pos_emb_adapter(q, k, cos, sin)
        self.assertEqual(result_q.shape, q.shape)
        self.assertEqual(result_k.shape, k.shape)
    
    def test_rotate_half(self):
        import adapter
        x = torch.randn(2, 4, 8, 16)
        result = adapter._rotate_half(x)
        x1 = x[..., : x.shape[-1] // 2]
        x2 = x[..., x.shape[-1] // 2 :]
        expected = torch.cat((-x2, x1), dim=-1)
        self.assertTrue(torch.allclose(result, expected))
    
    @patch('adapter.HAS_NPU', False)
    def test_apply_rotary_pos_emb_with_position_ids(self):
        import adapter
        q = torch.randn(2, 4, 8, 16)
        k = torch.randn(2, 4, 8, 16)
        cos = torch.randn(2, 8, 16)
        sin = torch.randn(2, 8, 16)
        position_ids = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7], [0, 1, 2, 3, 4, 5, 6, 7]])
        result_q, result_k = adapter._apply_rotary_pos_emb(q, k, cos, sin, position_ids)
        self.assertEqual(result_q.shape, q.shape)
        self.assertEqual(result_k.shape, k.shape)
    
    @patch('adapter.HAS_NPU', False)
    def test_apply_rotary_pos_emb_different_unsqueeze_dim(self):
        import adapter
        q = torch.randn(2, 8, 4, 16)
        k = torch.randn(2, 8, 4, 16)
        cos = torch.randn(2, 8, 16)
        sin = torch.randn(2, 8, 16)
        result_q, result_k = adapter._apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=2)
        self.assertEqual(result_q.shape, q.shape)
        self.assertEqual(result_k.shape, k.shape)
    
    @patch('adapter.HAS_NPU', False)
    def test_adapter_gelu_edge_cases(self):
        import adapter
        x = torch.tensor([0.0, float('inf'), -float('inf')])
        result = adapter.adapter_gelu(x)
        self.assertEqual(result[0], 0.0)
        self.assertEqual(result[1], float('inf'))
        self.assertTrue(torch.isnan(result[2]))
    
    @patch('adapter.HAS_NPU', False)
    def test_adpater_rmsnorm_different_dtypes(self):
        import adapter
        for dtype in [torch.float32, torch.float16, torch.bfloat16]:
            x = torch.randn(4, 8, dtype=dtype)
            weight = torch.randn(8, dtype=dtype)
            eps = 1e-5
            result = adapter.adpater_rmsnorm(x, weight, eps)
            self.assertEqual(result.dtype, dtype)
            self.assertEqual(result.shape, x.shape)
    
    @patch('adapter.HAS_NPU', False)
    def test_contiguous_for_channels_last_3d_different_shapes(self):
        import adapter
        shapes = [(2, 3, 4, 5, 6), (1, 3, 64, 64, 64), (4, 3, 224, 224, 224)]
        for shape in shapes:
            tensor = torch.randn(*shape)
            result = adapter.contiguous_for_channels_last_3d_memory_format(tensor)
            self.assertEqual(result.shape, shape)


class TestAdapterIntegration(unittest.TestCase):
    
    def setUp(self):
        try:
            import adapter
            self.adapter = adapter
        except ImportError:
            self.skipTest("adapter module not found")
    
    @patch('adapter.HAS_NPU', False)
    def test_full_workflow_gpu(self):
        import adapter
        x = torch.randn(2, 4, 8, 16)
        weight = torch.randn(16)
        eps = 1e-5
        
        gelu_result = adapter.adapter_gelu(x)
        norm_result = adapter.adpater_rmsnorm(x, weight, eps)
        
        x_5d = torch.randn(2, 3, 4, 8, 16)
        contiguous_result = adapter.contiguous_for_channels_last_3d_memory_format(x_5d)
        
        self.assertEqual(gelu_result.shape, x.shape)
        self.assertEqual(norm_result.shape, x.shape)
        self.assertEqual(contiguous_result.shape, x_5d.shape)
    
    @patch('adapter.HAS_NPU', False)
    def test_rotary_embedding_workflow(self):
        import adapter
        batch_size, heads, seq_len, head_dim = 2, 4, 8, 16
        q = torch.randn(batch_size, heads, seq_len, head_dim)
        k = torch.randn(batch_size, heads, seq_len, head_dim)
        cos = torch.randn(batch_size, seq_len, head_dim)
        sin = torch.randn(batch_size, seq_len, head_dim)
        
        q_rotated, k_rotated = adapter.apply_rotary_pos_emb_adapter(q, k, cos, sin)
        
        self.assertEqual(q_rotated.shape, q.shape)
        self.assertEqual(k_rotated.shape, k.shape)
        
        q_norm = torch.norm(q_rotated - q)
        k_norm = torch.norm(k_rotated - k)
        self.assertGreater(q_norm, 0)
        self.assertGreater(k_norm, 0)


if __name__ == '__main__':
    unittest.main()
