import triton
import triton.language as tl
from typing import Optional
from utils.benchmark_utils import (
    get_model_configs,
    get_available_models,
    torch_to_tl_dtype,
)
from op_tests.triton_tests.test_moe import input_helper, input_helper_int4_w4a16
from aiter import per_token_quant_hip, per_block_quant_wrapper
import torch
import argparse
from aiter.ops.triton.moe_op import fused_moe as triton_moe
import sys


def model_benchmark_configs(args):
    config_file = args.model_configs
    configs = get_model_configs(
        config_path=config_file, models="mistral" if args.model is None else args.model
    )
    moe_configs = []
    M = args.M if args.M else 4096  # check size
    # M, K, N, E, top_k

    for model_name, config in configs.items():
        N1 = config["intermediate_size"]
        K1 = config["hidden_size"]

        N2 = config["hidden_size"]
        K2 = config["intermediate_size"] // 2

        E = 8
        top_k = 2

        moe_configs.append((model_name, M, N1, K1, E, top_k))
        moe_configs.append((model_name, M, N2, K2, E, top_k))

    return moe_configs

moe_perf_model_cases_list = [
    #"M", "N", "K", "E", "top_k"
    (1,        256, 7168, 256, 8),
    (2,        256, 7168, 256, 8),
    (3,        256, 7168, 256, 8),
    (4,        256, 7168, 256, 8),
    # (5,        256, 7168, 256, 8),
    # (6,        256, 7168, 256, 8),
    # (7,        256, 7168, 256, 8),
    # (8,        256, 7168, 256, 8),
    # (9,        256, 7168, 256, 8),
    # (10,       256, 7168, 256, 8),
    # (11,       256, 7168, 256, 8),
    # (12,       256, 7168, 256, 8),
    # (13,       256, 7168, 256, 8),
    # (14,       256, 7168, 256, 8),
    # (15,       256, 7168, 256, 8),
    # (16,       256, 7168, 256, 8),
    # (17,       256, 7168, 256, 8),
    # (18,       256, 7168, 256, 8),
    # (19,       256, 7168, 256, 8),
    # (20,       256, 7168, 256, 8),
    # (21,       256, 7168, 256, 8),
    # (22,       256, 7168, 256, 8),
    # (23,       256, 7168, 256, 8),
    # (24,       256, 7168, 256, 8),
    # (25,       256, 7168, 256, 8),
    # (26,       256, 7168, 256, 8),
    # (27,       256, 7168, 256, 8),
    # (28,       256, 7168, 256, 8),
    # (29,       256, 7168, 256, 8),
    # (30,       256, 7168, 256, 8),
    # (31,       256, 7168, 256, 8),
    # (32,       256, 7168, 256, 8),
    # (64,       256, 7168, 256, 8),
    # (128,      256, 7168, 256, 8),
    # (256,      256, 7168, 256, 8),
    # (512,      256, 7168, 256, 8),
    # (1024,     256, 7168, 256, 8),
    # (2048,     256, 7168, 256, 8),
    # (4096,     256, 7168, 256, 8),
    # (8192,     256, 7168, 256, 8),
    # (16384,    256, 7168, 256, 8),
    # (32768,    256, 7168, 256, 8)
]

def run_benchmark(args):
    routed_weight = args.routed_weight
    int8_w8a16 = args.int8_w8a16
    fp8_w8a8 = args.fp8_w8a8
    int8_w8a8 = args.int8_w8a8
    int4_w4a16 = args.int4_w4a16
    int4_w4a8 = args.int4_w4a8
    group_size = args.group_size
    has_zp = args.has_zp
    dtype = arg_to_torch_dtype[args.dtype]
    fp8_type = arg_to_torch_dtype[args.fp8_type]

    if (fp8_w8a8 or int8_w8a8) and group_size is None:
       group_size = [128, 128]

    if (int4_w4a16 or int4_w4a8) and group_size is None:
       group_size = [0, 64]

    kernel_name = "_fused_moe_kernel"
    if (int8_w8a16 or int4_w4a16 or int4_w4a8) and (group_size is not None) and group_size[1] > 0:
        kernel_name = "_fused_moe_kernel_gptq_awq"
    quantype = "fp16" if dtype == torch.float16 else "bf16"
    if int4_w4a16:
        quantype = "int4_w4a16"
    elif int4_w4a8:
        quantype = "int4_w4a8"
    elif int8_w8a16:
        quantype = "int8_w8a16"
    elif int8_w8a8:
        quantype = "int8_w8a8"
    elif fp8_w8a8:
        quantype = "fp8_w8a8"
    kernel_name = quantype +  ("_w2_weight" if routed_weight else "_w1_weight") + kernel_name
    # x_vals_list = model_benchmark_configs(args)
    x_vals_list = moe_perf_model_cases_list
    x_names = ["M", "N", "K", "E", "top_k", "dtype", "group_size"]

    # line_names = ["Time (ms)", "TFLOPS", "Bandwidth (GB/s)"]
    # line_vals = ["time", "tflops", "bandwidth"]
    line_names = ["Time (ms)"]
    line_vals = ["time"]

    benchmark = triton.testing.Benchmark(
        x_names=x_names,
        x_vals=[t + (dtype, group_size, ) for t in x_vals_list],
        line_arg="metric",
        line_vals=line_vals,
        line_names=line_names,
        # styles=[("red", "-"), ("blue", "-"), ("yellow", "-")],
        # ylabel="ms / TFLOPS / GB/s",
        styles=[("red", "-")],
        ylabel="ms",
        plot_name=f"{kernel_name}-benchmark",
        args={},
    )

    @triton.testing.perf_report([benchmark])
    def bench_moe_gemm(M, N, K, E, top_k, group_size, metric, dtype=torch.float16):
        # # (M, K) * (top_k, N, K) -> (M, top_k, N). 2 for multiplication and accumulation
        # flops = 2.0 * M * top_k * K * N
        # # The weight is applied on the gemm product which has the shape of (M, top_k, N)
        # if routed_weight:
        #     flops += M * top_k * N

        # if fp8_w8a8:
        #     a_bytes = b_bytes = torch.tensor([], dtype=fp8_type).element_size()
        #     c_bytes = torch.tensor([], dtype=dtype).element_size()
        # elif int8_w8a16:
        #     b_bytes = torch.tensor([], dtype=torch.int8).element_size()
        #     a_bytes = c_bytes = torch.tensor([], dtype=dtype).element_size()
        # elif int4_w4a16:
        #     b_bytes = torch.tensor([], dtype=torch.int8).element_size() / 2
        #     a_bytes = c_bytes = torch.tensor([], dtype=dtype).element_size()
        # else:
        #     a_bytes = b_bytes = c_bytes = torch.tensor([], dtype=dtype).element_size()
        # TODO add the int4 case

        # (M, K) memory load for A (E,  N,  K) for B not (top_k,  N,  K) because we are in total bringing in all expert matrices into the chip from memory. It's just that not all multiply the same A.
        # mem_read = (M * K) * a_bytes + (E * N * K) * b_bytes

        # mem_write = (M * top_k * N) * c_bytes
        # mem = mem_read + mem_write
        if int4_w4a16 or int4_w4a8:
            assert group_size is not None and group_size[1] > 0
            (
                a,
                b,
                triton_out,
                _,
                b_zp,
                b_scale,
                topk_weights,
                topk_ids,
                sorted_token_ids,
                expert_ids,
                num_tokens_post_padded,
                config,
            ) = input_helper_int4_w4a16(
                M,
                N,
                K,
                top_k,
                E,
                routed_weight=routed_weight,
                dtype=dtype,
                group_size=group_size[1],
                has_zp=has_zp,
                use_int4_w4a8=int4_w4a8,
            )
            a_scale = None
            if int4_w4a8:
                a, a_scale = per_block_quant_wrapper((1, group_size[1]))(per_token_quant_hip)(a)
        else:
            (
                a,
                b,
                triton_out,
                _,
                b_zp,
                a_scale,
                b_scale,
                topk_weights,
                topk_ids,
                sorted_token_ids,
                expert_ids,
                num_tokens_post_padded,
                config,
            ) = input_helper(
                M,
                N,
                K,
                top_k,
                E,
                routed_weight=routed_weight,
                dtype=dtype,
                fp8_w8a8=fp8_w8a8,
                int8_w8a16=int8_w8a16,
                int8_w8a8=int8_w8a8,
                block_shape=group_size,
            )

        fn = lambda: triton_moe(
            a,
            b,
            triton_out,
            a_scale,
            b_scale,
            b_zp,
            topk_weights,
            topk_ids,
            sorted_token_ids,
            expert_ids,
            num_tokens_post_padded,
            routed_weight,
            top_k,
            torch_to_tl_dtype[dtype],
            use_fp8_w8a8=fp8_w8a8,
            use_int8_w8a8=int8_w8a8,
            use_int8_w8a16=int8_w8a16,
            use_int4_w4a16=int4_w4a16,
            use_int4_w4a8=int4_w4a8,
            block_shape=group_size,
            config=config,
        )

        ms = triton.testing.do_bench(fn, warmup=25, rep=100)

        # bandwidth = mem / (ms * 1e-3) * 1e-9  # GB/s
        # tflops = flops / ms * 1e-9

        # Return exactly one scalar depending on which metric is active
        if metric == "time":
            return ms
        # elif metric == "tflops":
        #     return tflops
        # elif metric == "bandwidth":
        #     return bandwidth
        else:
            raise ValueError("Unknown metric: " + metric)

    bench_moe_gemm.run(save_path=".", print_data=True)


def parse_args():
    parser = argparse.ArgumentParser(
        prog="Benchmark MoE GEMM",
        allow_abbrev=False,
    )
    # parser.add_argument(
    #     "-model_configs",
    #     type=str,
    #     default="utils/model_configs.json",
    #     help="Model config json file.",
    # )
    # available_models = get_available_models()  # Dynamically load model names
    # model_help = (
    #     "Model name to benchmark. Select from: ["
    #     + ", ".join(available_models)
    #     + "]. Use 'all' to benchmark all models or leave blank for the default benchmark script."
    # )
    # parser.add_argument("-model", type=str, default=None, help=model_help)
    # parser.add_argument("-M", type=int, default=0, help="M dimension")
    parser.add_argument(
        "-group_size", type=int, nargs=2, default=None, help="group_size for in4/int8 block quant, "
        "2 int or None, 'e.g., --group_size 128 128'"
    )
    parser.add_argument("-routed_weight", action="store_true", default=False)
    parser.add_argument("-int8_w8a16", action="store_true", default=False)
    parser.add_argument("-fp8_w8a8", action="store_true", default=False)
    parser.add_argument("-int4_w4a16", action="store_true", default=False)
    parser.add_argument("-int4_w4a8", action="store_true", default=False)
    parser.add_argument("-int8_w8a8", action="store_true", default=False)
    parser.add_argument("-has_zp", action="store_true", default=True)
    parser.add_argument("-dtype", default="fp16")
    parser.add_argument("-fp8_type", default="e4m3fnuz")
    args = parser.parse_args()
    return args

arg_to_torch_dtype = {
    "fp16": torch.float16,
    "bf16": torch.bfloat16,
    "fp32": torch.float32,
    "e5m2fnuz": torch.float8_e5m2fnuz,
    "e4m3fnuz": torch.float8_e4m3fnuz,
}

def main():
    args = parse_args()
    run_benchmark(args)


if __name__ == "__main__":
    sys.exit(main())
