import torch
import unittest
import torch_mlu_ops as ops
import random
from common_utils import *
from itertools import product
import copy
from typing import Optional

class TestMoeCastGating(BtTestCase):
    def op_impl_base(self, *args):
      input, weight = args
      input = input.to(torch.float)
      output = torch.matmul(input, weight.permute(1, 0))
      return output

    @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "moe_cast_gating not support MLU3XX device")
    def test_moe_cast_gating_random(self):
       for _ in range(1000):
          total_seq = random.randint(1, 32768)
          hidden_size = random.randint(1, 16384)
          expert_num = random.randint(1, 128)
          input_dtype_list = [torch.half]
          if torch_mlu.mlu.is_bf16_supported():
            input_dtype_list.append(torch.bfloat16)
          input_dtype = random.choice(input_dtype_list)
          weight_dtype = torch.float
          print("total_seqlen={}, hidden_size={}, expert_num={}, input_dtype={}, testing...".format(
                      total_seq, hidden_size, expert_num, input_dtype))
          input = torch.randn(total_seq, hidden_size, dtype=input_dtype, device="mlu")
          weight = torch.randn(expert_num, hidden_size, dtype=weight_dtype, device="mlu")

          tmo_out = ops.moe_cast_gating(input, weight)
          torch_out = self.op_impl_base(input, weight)

          self.assertTensorsEqual(tmo_out.cpu().float(), torch_out.cpu().float(), 1e-4,
                                  use_MSE=True, use_RAE=True)

    @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "moe_cast_gating not support MLU3XX device")
    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_prevent(self):
       func = ops.moe_cast_gating
       input = torch.randn(1024, 8192, dtype=torch.half, device="mlu")
       input = input.as_strided(input.shape, (100, 1))
       weight = torch.randn(128, 8192, dtype=torch.float, device="mlu")
       self.assertException("input must be contiguous", func, input, weight)
       input = input.contiguous()
       weight = weight.as_strided(weight.shape, (100, 1))
       self.assertException("weight must be contiguous", func, input, weight)
       weight = weight.contiguous()
       weight = weight.reshape(1, 128, 8192)
       self.assertException("weight.dim() == 2", func, input, weight)
       weight = torch.randn(128, 2048, dtype=torch.float, device="mlu")
       self.assertException("input.size(-1) == weight.size(-1)", func, input, weight)
       weight = torch.randn(128, 8192, dtype=torch.half, device="mlu")
       self.assertException("weight type need be torch::kFloat32", func, input, weight)
       weight = weight.to(torch.float)
       input = input.to(torch.float)
       self.assertException("input type need be torch::kFloat16 or torch::kBFloat16", func, input, weight)
       input = torch.randn(1024, 16388, dtype=torch.half, device="mlu")
       weight = torch.randn(128, 16388, dtype=torch.float, device="mlu")
       self.assertException("hidden_size > 0 && hidden_size <= 16384", func, input, weight)
       input = torch.randn(1024, 16384, dtype=torch.half, device="mlu")
       weight = torch.randn(129, 16384, dtype=torch.float, device="mlu")
       self.assertException("expert_num > 0 && expert_num <= 128", func, input, weight)

    @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "moe_cast_gating not support MLU3XX device")
    def test_inductor(self):
      m, hidden_size, expert_num, input_dtype = 1024, 4096, 32, torch.half
      input = torch.randn(m, hidden_size, dtype=input_dtype, device="mlu")
      weight = torch.randn(expert_num, hidden_size, dtype=torch.float, device="mlu")
      args = (input, weight)
      self.base_opcheck(torch.ops.torch_mlu_ops.moe_cast_gating, args)

if __name__ == "__main__":
  if "MLU3" not in torch.mlu.get_device_name():
    exit(run_unittest(TestMoeCastGating))
