import torch
import torch_mlu
import torch_mlu_ops as tmo
from common import benchmark_forward, save_to_csv
import argparse
from tabulate import tabulate
import os

e2e_time_param_dict_list = [
                            {"batch": 1, "seq_len": 1, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
                             "gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
                             "topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},

                            {"batch": 16, "seq_len": 1, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
                             "gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
                             "topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},

                            {"batch": 72, "seq_len": 1, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
                             "gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
                             "topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},

                            {"batch": 128, "seq_len": 1, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
                             "gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
                             "topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},

                            {"batch": 490, "seq_len": 1, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
                             "gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
                             "topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},

                            {"batch": 525, "seq_len": 1, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
                             "gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
                             "topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},

                            {"batch": 1, "seq_len": 1024, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
                             "gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
                             "topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},

                            {"batch": 1, "seq_len": 2048, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
                             "gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
                             "topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},

                            {"batch": 1, "seq_len": 4096, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
                             "gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
                             "topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},

                            {"batch": 1, "seq_len": 8192, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
                             "gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
                             "topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},

                            {"batch": 1, "seq_len": 32768, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
                             "gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
                             "topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},
                            ]

def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('--repeat_times', type=int, default=10, help='repeat times for testing')
  parser.add_argument('--csv', action='store_true', help='write the report data to csv')
  parser.add_argument('-o', type=str, help='specify the output folder name under --csv mode')

  args = parser.parse_args()
  device = 'mlu'

  titles = ["batch", "seq_len", "hidden_size", "inner_size", "gated_ffn", "num_expert", "topk", "act_mode", "quant_weight", "dtype", "hardware_time(us)", "e2e_latency(us)"]
  contents = []
  for params_dict in e2e_time_param_dict_list:
    batch = params_dict["batch"]
    seq_len = params_dict["seq_len"]
    hidden_size = params_dict["hidden_size"]
    inner_size = params_dict["inner_size"]
    gated_ffn = params_dict["gated_ffn"]
    act_mode = params_dict["act_mode"]
    num_expert = params_dict["num_expert"]
    start_expert_id = params_dict["start_expert_id"]
    expert_size = params_dict["expert_size"]
    topk = params_dict["topk"]
    has_residual = params_dict["has_residual"]
    smooth_quant = params_dict["smooth_quant"]
    renormalize = params_dict["renormalize"]
    input_dtype_list = params_dict["dtype"]
    # print(f"batch:{batch}, seq_len:{seq_len}, hidden_size:{hidden_size}, inner_size:{inner_size}, "
    #       f"gated_ffn:{gated_ffn}, act_mode:{act_mode}, num_expert:{num_expert}, topk:{topk}")
    for dtype in input_dtype_list:
      if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
        dtype = torch.half
      hidden_states = torch.randn(batch, seq_len, hidden_size, device=device, dtype=dtype)
      router_logit = torch.randn(batch, seq_len, num_expert, device=device, dtype=torch.float32)

      if False:  # print token_count
        softmax = torch.softmax(router_logit.view(-1, router_logit.size(-1)), dim=1)
        topk_logit, expert_id = torch.topk(softmax, k=topk, dim=1)
        if renormalize:
            topk_logit = topk_logit / topk_logit.sum(-1).unsqueeze(1)
        sorted_expert_id, indices = expert_id.int().flatten().sort()
        token_cout = torch.bincount(sorted_expert_id, minlength=num_expert).int()
        print(token_cout)

      residual = None
      if has_residual:
        residual = torch.randn(batch, seq_len, hidden_size, device=device, dtype=dtype)
      weight1 = torch.randn(num_expert, inner_size*(1+gated_ffn), hidden_size, device=device, dtype=dtype)
      bias1 = None # torch.randn(expert_num, inner_size*(1+gated), device=device, dtype=data_type)
      weight2 = torch.randn(num_expert, hidden_size, inner_size, device=device, dtype=dtype)
      bias2 = None # torch.randn(expert_num, hidden_size, device=device, dtype=data_type)
      input_smooth, act_smooth, w1_scale, w2_scale = None, None, None, None
      if smooth_quant:
          input_smooth = torch.randn(expert_size, hidden_size, device=device, dtype=torch.float32).abs() + 0.1
          act_smooth = torch.randn(expert_size, inner_size, device=device, dtype=torch.float32).abs() + 0.1
          weight1 = torch.randint(-128, 127, (num_expert, inner_size*(1+gated_ffn), hidden_size)).to(torch.int8).mlu()
          weight2 = torch.randint(-128, 127, (num_expert, hidden_size, inner_size)).to(torch.int8).mlu()
          w1_scale = torch.randn(expert_size, (1+gated_ffn)*inner_size).to(device).to(torch.float32)
          w2_scale = torch.randn(expert_size, hidden_size).to(device).to(torch.float32)

      hardware_time, e2e_time = benchmark_forward(tmo.fused_moe,
                                                  hidden_states,
                                                  router_logit,
                                                  weight1[start_expert_id:start_expert_id+expert_size],
                                                  weight2[start_expert_id:start_expert_id+expert_size],
                                                  bias1,
                                                  bias2,
                                                  residual,
                                                  input_smooth,
                                                  act_smooth,
                                                  w1_scale,
                                                  w2_scale,
                                                  topk,
                                                  renormalize,
                                                  gated_ffn,
                                                  act_mode,
                                                  start_expert_id,
                                                  repeats=args.repeat_times)
      content = [f"{batch}", f"{seq_len}", f"{hidden_size}", f"{inner_size}", f"{gated_ffn}", f"{num_expert}", f"{topk}", f"{act_mode}", f"{smooth_quant}", f"{dtype}", f"{hardware_time}", f"{e2e_time}"]
      contents.append(content)
  table = [titles] + contents
  print(tabulate(table, headers="firstrow", tablefmt="grid"))

  if args.csv:
    current_file_path = __file__
    _, file_name = os.path.split(current_file_path)
    save_to_csv(table, args.o, file_name)

if __name__=="__main__":
  main()
