# ruff: noqa
"""
Tune fp8_index configs by benchmarking. M 1~4096, N 64~128K.
Generates M_REPR_TABLE, N_REPR_TABLE and config lookup for fp8_index.py.
"""
import sys
import os
import json
import argparse
import torch
import numpy as np
from concurrent.futures import ProcessPoolExecutor
from typing import Optional, List, Tuple, Dict
from tilelang.profiler import do_bench_cudagraph
from aiter.ops.tilelang.fp8_index import act_quant, _get_fp8_index_kernel

# Add parent directory to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../tilelang_tests'))
from test_fp8_index import ref_act_quant, ref_fp8_index


def _bench_one_point_worker(args: Tuple) -> Tuple[Tuple[int, int], Tuple[int, int, int]]:
    """Worker: benchmark one (m, n) on given device. Returns ((m, n), (m_split, blk_n1, blk_n2))."""
    m, n, b, h, d, compiled, device_id = args
    torch.cuda.set_device(device_id)
    m_split, blk_n1, blk_n2, _ = tune_one_point(
        m, n, b=b, h=h, d=d, quiet=True,
        skip_compile=True, compiled_configs=compiled, device_id=device_id,
    )
    return ((m, n), (m_split, blk_n1, blk_n2))


def _compile_kernels_worker(args: Tuple) -> List[Tuple[int, int, int]]:
    """Worker: compile kernels only (no run, pure host). Returns configs that compiled successfully."""
    configs, h, d = args

    valid = []
    for m_split, blk_n1, blk_n2 in configs:
        try:
            _get_fp8_index_kernel(h, d, m_split, blk_n1, blk_n2, threads=256, clear_accum=False)
            valid.append((m_split, blk_n1, blk_n2))
        except Exception:
            pass
    return valid


def get_m_repr_points() -> List[int]:
    """M in [1, 4096]. 1~128: 2^n; 128~512: step 128; 512~1024: step 128; 1024~2048: step 256; 2048~4096: step 512.
    Use < so x flows to next segment; last segment uses <= to include 4096."""
    points = []
    x = 1
    while x < 128:
        points.append(x)
        x *= 2
    while x < 512:
        points.append(x)
        x += 128
    while x < 1024:
        points.append(x)
        x += 128
    while x < 2048:
        points.append(x)
        x += 256
    while x <= 4096:
        points.append(x)
        x += 512
    return points


def get_n_repr_points() -> List[int]:
    """N in [64, 131072]. Segments: 64~512: 2^n; 512~1024: step 256; 1024~2048: step 512;
    2048~4096: step 1024; 4096~8192: step 2048; 8192~16384: step 4096; 16384~131072: step 8192.
    Use < so x flows to next segment; last segment uses <= to include 131072."""
    points = []
    x = 64
    while x < 512:
        points.append(x)
        x *= 2
    while x < 1024:
        points.append(x)
        x += 256
    while x < 2048:
        points.append(x)
        x += 512
    while x < 4096:
        points.append(x)
        x += 1024
    while x < 8192:
        points.append(x)
        x += 2048
    while x < 16384:
        points.append(x)
        x += 4096
    while x <= 131072:
        points.append(x)
        x += 8192
    return points


def _enumerate_compute_configs(b: int, m: int, n: int, h: int, d: int) -> List[Tuple[int, int, int]]:
    """Enumerate (m_split, blk_n1, blk_n2) candidates. Pruned for tuning."""
    cu_count = torch.cuda.get_device_properties("cuda").multi_processor_count

    def next_pow2(x):
        if x <= 0:
            return 1
        return 1 << (x - 1).bit_length()

    ms_candidates = []
    for ms in [8, 4, 2, 1]:
        if ms > m:
            continue
        if ms * h > 256:
            continue
        if ms == 1 and (m // 2) * (n // 64) > 8 * cu_count:
            continue
        if ms == 2 and (m // 4) * (n // 64) > 8 * cu_count:
            continue
        ms_candidates.append(ms)
    if not ms_candidates:
        ms_candidates = [1]

    nblock_upper = min(8192, next_pow2(n))
    nblock_lower = max(64, next_pow2(n // (4 * cu_count)))
    nblock_candidates = [nblock_lower]
    x = nblock_lower * 2
    while x <= nblock_upper:
        nblock_candidates.append(x)
        x *= 2

    configs = []
    for ms in ms_candidates:
        for blk_n1 in nblock_candidates:
            if blk_n1 > n:
                continue
            if ms * h > 128:
                blk_n2 = 64
            else:
                blk_n2 = 128 if blk_n1 >= 512 else 64
            if blk_n1 % blk_n2 == 0:
                configs.append((ms, blk_n1, blk_n2))
    return configs


def _collect_all_configs(
    m_points: List[int],
    n_points: List[int],
    b: int,
    h: int,
    d: int,
) -> List[Tuple[int, int, int]]:
    """Collect all unique (m_split, blk_n1, blk_n2) across all (m, n) points."""
    all_configs = set()
    for m in m_points:
        for n in n_points:
            if n < 64:
                continue
            configs = _enumerate_compute_configs(b, m, n, h, d)
            all_configs.update(configs)
    return sorted(all_configs)


def _compile_all_kernels(
    configs: List[Tuple[int, int, int]],
    h: int,
    d: int,
    workers: int,
) -> set:
    """Compile all kernels in batch. Returns set of successfully compiled (m_split, blk_n1, blk_n2)."""
    if not configs:
        return set()
    compiled = set()
    if workers <= 1:
        compiled.update(_compile_kernels_worker((configs, h, d)))
    else:
        chunks = [[] for _ in range(workers)]
        for i, cfg in enumerate(configs):
            chunks[i % workers].append(cfg)
        task_args = [(chunk, h, d) for chunk in chunks if chunk]
        with ProcessPoolExecutor(max_workers=workers) as ex:
            for result in ex.map(_compile_kernels_worker, task_args):
                compiled.update(result)
    torch.cuda.synchronize()
    return compiled


def tune_one_point(
    m: int,
    n: int,
    b: int = 1,
    h: int = 32,
    d: int = 128,
    block_size: int = 128,
    scale_fmt=None,
    dtype=torch.bfloat16,
    quiet: bool = True,
    workers: int = 1,
    device_id: int = 0,
    skip_compile: bool = False,
    compiled_configs: Optional[set] = None,
) -> Tuple[int, int, int, float]:
    """Tune for (m, n) and return best (m_split, blk_n1, blk_n2), best_ms.
    If skip_compile=True, only benchmark (kernels must be pre-compiled). compiled_configs filters valid configs."""
    torch.random.manual_seed(0)
    device = "cuda"

    q = torch.randn((b, m, h, d), dtype=dtype, device=device)
    k = torch.randn((b, n, d), dtype=dtype, device=device)
    q_fp8, q_scale = act_quant(q.contiguous(), block_size=block_size, scale_fmt=scale_fmt)
    k_fp8, k_scale = act_quant(k.contiguous(), block_size=block_size, scale_fmt=scale_fmt)
    k_scale_index = k_scale[..., 0]
    q_scale_first = q_scale[..., 0]
    weights_base = torch.randn((b, m, h), dtype=torch.float32, device=device) * 0.1
    softmax_scale = 1.0 / np.sqrt(h)
    q_s = weights_base * q_scale_first * softmax_scale

    q_fp8 = q_fp8.contiguous()
    k_fp8 = k_fp8.contiguous()
    q_s = q_s.contiguous()
    k_scale_index = k_scale_index.contiguous()

    configs = _enumerate_compute_configs(b, m, n, h, d)
    if skip_compile and compiled_configs is not None:
        configs = [c for c in configs if c in compiled_configs]
        if not configs:
            skip_compile = False
            configs = _enumerate_compute_configs(b, m, n, h, d)

    import builtins
    _orig_print = builtins.print
    if quiet:
        def _noop(*args, **kwargs):
            if args and "[fp8_index] kernel config" in str(args[0]):
                return
            _orig_print(*args, **kwargs)
        builtins.print = _noop

    # Phase 1: compile kernels (skip if skip_compile)
    valid_configs = []
    if not skip_compile:
        if workers <= 1:
            for m_split, blk_n1, blk_n2 in configs:
                try:
                    kernel = _get_fp8_index_kernel(h, d, m_split, blk_n1, blk_n2, threads=256, clear_accum=False)
                    kernel(q_fp8, q_s, k_fp8, k_scale_index)
                    valid_configs.append((m_split, blk_n1, blk_n2))
                except Exception:
                    pass
        else:
            chunks = [[] for _ in range(workers)]
            for i, cfg in enumerate(configs):
                chunks[i % workers].append(cfg)
            task_args = [(chunk, h, d) for chunk in chunks if chunk]
            with ProcessPoolExecutor(max_workers=workers) as ex:
                for result in ex.map(_compile_kernels_worker, task_args):
                    valid_configs.extend(result)
        torch.cuda.synchronize()
    else:
        valid_configs = list(configs)

    # Phase 2: trial run to filter failing configs, then benchmark
    best_ms, best_config = float("inf"), (1, 128, 64)
    for m_split, blk_n1, blk_n2 in valid_configs:
        try:
            kernel = _get_fp8_index_kernel(h, d, m_split, blk_n1, blk_n2, threads=256, clear_accum=False)
            kernel(q_fp8, q_s, k_fp8, k_scale_index)  # trial run, drop if fails
        except Exception:
            continue

        try:
            def fn():
                return kernel(q_fp8, q_s, k_fp8, k_scale_index)

            ms = do_bench_cudagraph(fn)
            if ms < best_ms:
                best_ms, best_config = ms, (m_split, blk_n1, blk_n2)
        except Exception:
            pass

    builtins.print = _orig_print
    return best_config[0], best_config[1], best_config[2], best_ms


def run_tuning(
    m_points: Optional[List[int]] = None,
    n_points: Optional[List[int]] = None,
    b: int = 1,
    h: int = 32,
    d: int = 128,
    output_path: Optional[str] = None,
    workers: int = 1,
    bench_workers: int = 1,
    device_ids: Optional[List[int]] = None,
    compile_only: bool = False,
    bench_only: bool = False,
) -> Tuple[List[int], List[int], Dict[Tuple[int, int], Tuple[int, int, int]]]:
    """Run tuning and return M_REPR_TABLE, N_REPR_TABLE, config_map.
    Phase 1: compile all kernels. Phase 2: benchmark each shape.
    bench_workers: parallel bench workers, each on device_ids[i % len(device_ids)].
    compile_only: only compile, no bench. bench_only: only bench, assume kernels compiled."""
    m_points = m_points or get_m_repr_points()
    n_points = n_points or get_n_repr_points()

    all_configs = _collect_all_configs(m_points, n_points, b, h, d)
    compiled: set = set()

    # Phase 1: compile all kernels (skip if bench_only)
    if not bench_only:
        print(f"Phase 1: Compiling {len(all_configs)} unique kernels (workers={workers}) ...", flush=True)
        compiled = _compile_all_kernels(all_configs, h, d, workers)
        print(f"  -> {len(compiled)} kernels compiled successfully", flush=True)
        if compile_only:
            return m_points, n_points, {}

    # Phase 2: tune each (m, n) with pre-compiled kernels
    if bench_only:
        compiled = set(all_configs)  # assume all compiled
    device_ids = device_ids or [0]
    tasks = [(m, n) for m in m_points for n in n_points if n >= 64]
    config_map = {}

    if bench_workers <= 1:
        total = len(tasks)
        for idx, (m, n) in enumerate(tasks, 1):
            print(f"[{idx}/{total}] Tuning m={m} n={n} ...", flush=True)
            m_split, blk_n1, blk_n2, ms = tune_one_point(
                m, n, b=b, h=h, d=d, quiet=True,
                skip_compile=True, compiled_configs=compiled, device_id=device_ids[0],
            )
            config_map[(m, n)] = (m_split, blk_n1, blk_n2)
            print(f"  -> best: m_split={m_split} blk_n1={blk_n1} blk_n2={blk_n2}  ms={ms:.4f}")
    else:
        n_dev = len(device_ids)
        task_args = [
            (m, n, b, h, d, compiled, device_ids[i % n_dev])
            for i, (m, n) in enumerate(tasks)
        ]
        print(f"Phase 2: Benchmarking {len(tasks)} shapes (workers={bench_workers}, devices={device_ids}) ...", flush=True)
        with ProcessPoolExecutor(max_workers=bench_workers) as ex:
            for (m, n), cfg in ex.map(_bench_one_point_worker, task_args):
                config_map[(m, n)] = cfg
                print(f"  m={m} n={n} -> {cfg}", flush=True)

    return m_points, n_points, config_map


def emit_fp8_index_config(
    m_repr: List[int],
    n_repr: List[int],
    config_map: Dict[Tuple[int, int], Tuple[int, int, int]],
    output_path: str,
) -> None:
    """Emit Python config for fp8_index.py."""
    lines = [
        "# Auto-generated by tune_fp8_index.py. Do not edit manually.",
        "",
        "from typing import Tuple, Dict",
        "import bisect",
        "",
        "M_REPR_TABLE = " + repr(m_repr),
        "N_REPR_TABLE = " + repr(n_repr),
        "",
        "CONFIG_MAP: Dict[Tuple[int, int], Tuple[int, int, int]] = {",
    ]
    for (m, n), (ms, b1, b2) in sorted(config_map.items()):
        lines.append(f"    ({m}, {n}): ({ms}, {b1}, {b2}),")
    lines.append("}")
    lines.append("")
    lines.append("")
    lines.append("def get_tuned_config(m: int, n: int) -> Tuple[int, int, int]:")
    lines.append('    """Lookup tuned config for (m, n) using floor. Returns (m_split, blk_n1, blk_n2)."""')
    lines.append("    if m <= M_REPR_TABLE[0]:")
    lines.append("        m_repr = M_REPR_TABLE[0]")
    lines.append("    elif m >= M_REPR_TABLE[-1]:")
    lines.append("        m_repr = M_REPR_TABLE[-1]")
    lines.append("    else:")
    lines.append("        idx = bisect.bisect_right(M_REPR_TABLE, m) - 1")
    lines.append("        m_repr = M_REPR_TABLE[idx]")
    lines.append("    if n <= N_REPR_TABLE[0]:")
    lines.append("        n_repr = N_REPR_TABLE[0]")
    lines.append("    elif n >= N_REPR_TABLE[-1]:")
    lines.append("        n_repr = N_REPR_TABLE[-1]")
    lines.append("    else:")
    lines.append("        idx = bisect.bisect_right(N_REPR_TABLE, n) - 1")
    lines.append("        n_repr = N_REPR_TABLE[idx]")
    lines.append("    return CONFIG_MAP[(m_repr, n_repr)]")
    lines.append("")

    out_dir = os.path.dirname(output_path)
    if out_dir:
        os.makedirs(out_dir, exist_ok=True)
    with open(output_path, "w") as f:
        f.write("\n".join(lines))
    print(f"Wrote {output_path}")


def main():
    parser = argparse.ArgumentParser(description="Tune fp8_index configs")
    parser.add_argument("--output", "-o", default=None, help="Output config path (default: aiter/ops/tilelang/fp8_index_tuned_config.py)")
    parser.add_argument("--workers", "-j", type=int, default=1, help="Parallel workers for kernel compilation (default: 1)")
    parser.add_argument("--bench-workers", type=int, default=1, help="Parallel workers for benchmark (default: 1)")
    parser.add_argument("--devices", type=str, default=None, help="Comma-separated device IDs for bench (e.g. 0,1,2,3), default: 0")
    parser.add_argument("--compile-only", action="store_true", help="Only compile kernels, no benchmark")
    parser.add_argument("--bench-only", action="store_true", help="Only benchmark, assume kernels already compiled")
    parser.add_argument("--m", type=str, default=None, help="Comma-separated M points (override)")
    parser.add_argument("--n", type=str, default=None, help="Comma-separated N points (override)")
    parser.add_argument("--b", type=int, default=1, help="Batch size (default: 1)")
    parser.add_argument("--h", type=int, default=32, help="Head count (default: 32)")
    parser.add_argument("--d", type=int, default=128, help="Head dimension (default: 128)")
    parser.add_argument("--json", action="store_true", help="Also emit JSON")
    args = parser.parse_args()

    device_ids = [int(x) for x in args.devices.split(",")] if args.devices else [0]
    torch.cuda.set_device(device_ids[0])

    m_points = [int(x) for x in args.m.split(",")] if args.m else None
    n_points = [int(x) for x in args.n.split(",")] if args.n else None

    output_path = args.output
    if output_path is None:
        # op_tests/op_benchmarks/tilelang/ -> aiter/
        script_dir = os.path.dirname(os.path.abspath(__file__))
        aiter_root = os.path.dirname(os.path.dirname(os.path.dirname(script_dir)))
        output_path = os.path.join(aiter_root, "aiter", "ops", "tilelang", "fp8_index_tuned_config.py")

    print("=" * 80)
    print(f"fp8_index tuning: M 1~4096, N 64~128K, b={args.b}, h={args.h}, d={args.d}")
    print("=" * 80)

    m_repr, n_repr, config_map = run_tuning(
        m_points=m_points,
        n_points=n_points,
        b=args.b,
        h=args.h,
        d=args.d,
        workers=args.workers,
        bench_workers=args.bench_workers,
        device_ids=device_ids,
        compile_only=args.compile_only,
        bench_only=args.bench_only,
    )

    if not args.compile_only:
        emit_fp8_index_config(m_repr, n_repr, config_map, output_path)

    if args.json and not args.compile_only:
        json_path = output_path.replace(".py", ".json")
        with open(json_path, "w") as f:
            json.dump(
                {"M_REPR_TABLE": m_repr, "N_REPR_TABLE": n_repr, "CONFIG_MAP": {f"{k[0]}_{k[1]}": v for k, v in config_map.items()}},
                f,
                indent=2,
            )
        print(f"Wrote {json_path}")

    print("=" * 80)
    if args.compile_only:
        print("Compile done. Run without --compile-only to benchmark and emit config.")
    else:
        print("Tuning done. fp8_index.py will use get_tuned_config() when the file exists.")
        print("To refine repr points: re-run with --m/--n to add points where configs change.")
    print("=" * 80)


if __name__ == "__main__":
    import multiprocessing
    try:
        multiprocessing.set_start_method("spawn")
    except RuntimeError:
        pass
    main()
