import argparse
import sys
import triton
from aiter.ops.triton.utils.types import torch_to_triton_dtype
import aiter.ops.triton.utils._triton.arch_info as arch_info
from aiter.ops.triton.moe.moe_op_mxfp4 import fused_moe_mxfp4
from aiter.ops.triton.moe.moe_op_mxfp4_silu_fused import fused_moe_mxfp4_silu
from op_tests.triton_tests.moe.test_moe_mx import input_helper
from op_tests.op_benchmarks.triton.utils.benchmark_utils import (
    get_available_models,
    get_model_configs,
    get_caller_name_no_ext,
    print_vgpr,
)


def model_benchmark_configs(args):
    config_file = args.model_configs
    configs = get_model_configs(
        config_path=config_file, models="mixtral" if args.model is None else args.model
    )
    moe_configs = []
    M = args.M if args.M else 128  # 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


def run_benchmark(args):
    routed_weight = args.routed_weight
    a_dtype_str = args.a_dtype
    b_dtype_str = "mxfp4_e2m1"
    swizzle_mx = args.swizzle_mx
    silu_fused = args.silu_fused
    print(f"MoE Benchmark {a_dtype_str} x {b_dtype_str}")

    x_vals_list = model_benchmark_configs(args)
    x_names = ["model", "M", "N", "K", "E", "top_k"]

    line_names = ["Time_(ms)", "TFLOPS", "Bandwidth_(GB/s)"]
    line_vals = ["time", "tflops", "bandwidth"]

    benchmark = triton.testing.Benchmark(
        x_names=x_names,
        x_vals=x_vals_list,
        line_arg="metric",
        line_vals=line_vals,
        line_names=line_names,
        styles=[("red", "-"), ("blue", "-"), ("yellow", "-")],
        ylabel="ms / TFLOPS / GB/s",
        plot_name=get_caller_name_no_ext(),
        args={"a_dtype": a_dtype_str, "swizzle_mx": swizzle_mx},
    )

    @triton.testing.perf_report([benchmark])
    def bench_moe_gemm(M, N, K, E, top_k, metric, a_dtype, swizzle_mx, model=None):

        (
            a_tri,
            b_tri,
            c_tri,
            c_tri_silu,
            a_scale,
            b_scale,
            a_mx_scales,
            b_mx_scales,
            topk_weights,
            topk_ids,
            sorted_token_ids,
            expert_ids,
            num_tokens_post_padded,
            top_k,
            config,
        ) = input_helper(M, N, K, top_k, E, a_dtype_str, b_dtype_str)
        # (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

        # Variables to compute bandwidth
        mem_read = (
            a_tri.numel() * a_tri.element_size() + b_tri.numel() * b_tri.element_size()
        )
        mem_write = c_tri.numel() * c_tri.element_size()

        if silu_fused:
            mem = mem_read + (mem_write // 2)
            flops += M * top_k * N
        else:
            mem = mem_read + mem_write

        mem = mem_read + mem_write

        fused_moe = fused_moe_mxfp4_silu if silu_fused else fused_moe_mxfp4
        output_tensor = c_tri_silu if silu_fused else c_tri

        fn = lambda: fused_moe(  # noqa: E731
            a_tri,
            b_tri,
            output_tensor,
            a_scale,
            b_scale,
            a_mx_scales,
            b_mx_scales,
            topk_weights,
            topk_ids,
            sorted_token_ids,
            expert_ids,
            num_tokens_post_padded,
            routed_weight,
            top_k,
            swizzle_mx,
            swizzle_mx,
            config,
            torch_to_triton_dtype[c_tri.dtype],
        )

        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="." if args.o else None, print_data=True)


def parse_args():
    parser = argparse.ArgumentParser(
        prog="Benchmark MoE with micro scaled format",
        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, help="M dimension")
    parser.add_argument("--routed-weight", action="store_true")
    parser.add_argument("--swizzle-mx", action="store_true")
    parser.add_argument("-silu_fused", action="store_true", default=False)
    parser.add_argument("-print_vgpr", action="store_true", default=False)
    parser.add_argument(
        "-A",
        "--a-dtype",
        type=str,
        choices=["bf16", "fp16", "fp8_e5m2", "mxfp4_e2m1"],
        default="mxfp4_e2m1",
    )
    parser.add_argument(
        "-o", action="store_true", help="Write performance results to CSV file"
    )
    args = parser.parse_args()
    return args


def main():
    if not (arch_info.is_fp4_avail()):
        print("MXFP4 not supported on this architecture")
        sys.exit(0)

    args = parse_args()
    if args.print_vgpr:
        print("Retrieving VGPR usage for Triton kernels...")
        fun = lambda: run_benchmark(args)  # noqa: E731
        print_vgpr(fun, get_caller_name_no_ext())
        return 0
    run_benchmark(args)


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