import argparse
import csv

import torch
import torch.nn.functional as F

import aiter
from aiter.test_common import run_perftest, checkAllclose
from aiter.ops.triton.fused_moe import (
    triton_silu_and_mul,
    triton_gelu_and_mul,
    triton_relu2,
    activation_and_mul_kernel,
)


def torch_silu_and_mul(input: torch.Tensor) -> torch.Tensor:
    d = input.shape[-1] // 2
    x, y = input.split([d, d], dim=-1)
    return F.silu(x) * y


def torch_gelu_and_mul(input: torch.Tensor) -> torch.Tensor:
    d = input.shape[-1] // 2
    x, y = input.split([d, d], dim=-1)
    return F.gelu(x, approximate="none") * y


def torch_relu2(x: torch.Tensor) -> torch.Tensor:
    return torch.square(F.relu(x))


def get_best_config(cache: dict, m: int, n: int, act: int):
    # Prefer exact key match (common in Triton autotune cache)
    for k, v in cache.items():
        if isinstance(k, tuple) and len(k) >= 3:
            if k[0] == m and k[1] == n and k[2] == act:
                return v
    # Fallback: match by membership if key structure differs
    for k, v in cache.items():
        if isinstance(k, tuple) and (m in k and n in k and act in k):
            return v
    return None


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--m_start", type=int, default=1)
    parser.add_argument("--m_end", type=int, default=65536)
    parser.add_argument("--n", type=int, default=704)
    parser.add_argument("--dtype", type=str, default="bf16", choices=["fp16", "bf16"])
    parser.add_argument("--num_iters", type=int, default=5)
    parser.add_argument("--num_warmup", type=int, default=1)
    parser.add_argument(
        "--out_csv",
        type=str,
        default="act_and_mul_m_1_65536_n_704.csv",
    )
    args = parser.parse_args()

    dtype = torch.float16 if args.dtype == "fp16" else torch.bfloat16
    device = "cuda"

    with open(args.out_csv, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow([
            "m",
            "n",
            "dtype",
            "act",
            "us",
            "TB/s",
            "err_ratio",
            "best_config",
        ])
        f.flush()

        m = args.m_start
        while m <= args.m_end:
            n = args.n
            inp = torch.randn(m, n, dtype=dtype, device=device)
            out = torch.empty((m, n // 2), dtype=dtype, device=device)

            # Correctness + perf for SiLU (compare against aiter)
            ref = torch.empty_like(out)
            aiter.silu_and_mul(ref, inp)
            triton_silu_and_mul(out, inp)
            err = checkAllclose(ref, out, rtol=1e-2, atol=1e-2, printLog=False)
            _, us = run_perftest(
                triton_silu_and_mul,
                out,
                inp,
                num_iters=args.num_iters,
                num_warmup=args.num_warmup,
            )
            tbps = (inp.nbytes + out.nbytes) / us / 1e6
            cfg = None #get_best_config(activation_and_mul_kernel.cache, m, n // 2, 0)
            writer.writerow([m, n, str(dtype), "silu", us, tbps, err, str(cfg)])

            # Correctness + perf for GELU (compare against aiter)
            ref = torch.empty_like(out)
            aiter.gelu_and_mul(ref, inp)
            triton_gelu_and_mul(out, inp)
            err = checkAllclose(ref, out, rtol=1e-2, atol=1e-2, printLog=False)
            _, us = run_perftest(
                triton_gelu_and_mul,
                out,
                inp,
                num_iters=args.num_iters,
                num_warmup=args.num_warmup,
            )
            tbps = (inp.nbytes + out.nbytes) / us / 1e6
            cfg = None #get_best_config(activation_and_mul_kernel.cache, m, n // 2, 1)
            writer.writerow([m, n, str(dtype), "gelu", us, tbps, err, str(cfg)])

            # Elementwise ReLU² (same shape in/out; no gate mul)
            inp_r = torch.randn(m, n, dtype=dtype, device=device)
            out_r = torch.empty_like(inp_r)
            ref_r = torch_relu2(inp_r)
            triton_relu2(out_r, inp_r)
            err = checkAllclose(ref_r, out_r, rtol=1e-2, atol=1e-2, printLog=False)
            _, us = run_perftest(
                triton_relu2,
                out_r,
                inp_r,
                num_iters=args.num_iters,
                num_warmup=args.num_warmup,
            )
            tbps = (inp_r.nbytes + out_r.nbytes) / us / 1e6
            cfg = None  # get_best_config(relu2_kernel.cache, m, n, ...)
            writer.writerow([m, n, str(dtype), "relu2", us, tbps, err, str(cfg)])

            if m % 256 == 0:
                f.flush()
                aiter.logger.info(f"progress: {m}/{args.m_end}")
            m *= 2

    aiter.logger.info(f"saved: {args.out_csv}")


if __name__ == "__main__":
    main()
