# SPDX-License-Identifier: MIT


import math
import os

import torch
from aiter import dtypes
from torch import Tensor

from ..jit.core import compile_ops
from ..jit.utils.chip_info import get_cu_num, get_gfx
from ..jit.utils.torch_guard import torch_compile_guard


def _truthy_env(name: str) -> bool:
    v = os.environ.get(name, "").strip().lower()
    return v in ("1", "true", "yes", "on")


def _round_to_tf32_like_tilekernels(x: torch.Tensor) -> torch.Tensor:
    return (x.view(torch.int32) + 0x1000).view(torch.float32)


@compile_ops("module_mhc")
def mhc_pre_gemm_sqrsum(
    out: Tensor,
    sqrsum: Tensor,
    x: Tensor,
    fn: Tensor,
    tile_k: int = 128,  # 64 or 128
    use_tf32: bool = False,
) -> None: ...


@compile_ops("module_mhc")
def mhc_pre_gemm_sqrsum_stage1_m128(
    out: Tensor,
    sqrsum: Tensor,
    x: Tensor,
    fn: Tensor,
    use_tf32: bool = False,
) -> None: ...


@compile_ops("module_mhc")
def mhc_pre_reduce_splitk(
    out_red: Tensor,
    sqrsum_red: Tensor,
    out: Tensor,
    sqrsum: Tensor,
) -> None: ...


@compile_ops("module_mhc")
def mhc_pre_big_fuse(
    post_mix: Tensor,
    comb_mix: Tensor,
    layer_input: Tensor,
    gemm_out_mul: Tensor,
    gemm_out_sqrsum: Tensor,
    hc_scale: Tensor,
    hc_base: Tensor,
    residual: Tensor,
    rms_eps: float = 1e-6,
    hc_pre_eps: float = 1e-6,
    hc_sinkhorn_eps: float = 1e-6,
    hc_post_mult_value: float = 1.0,
    sinkhorn_repeat: int = 20,
) -> None: ...


@compile_ops("module_mhc")
def mhc_pre_big_fuse_tlstyle(
    post_mix: Tensor,
    comb_mix: Tensor,
    layer_input: Tensor,
    gemm_out_mul: Tensor,
    gemm_out_sqrsum: Tensor,
    hc_scale: Tensor,
    hc_base: Tensor,
    residual: Tensor,
    rms_eps: float = 1e-6,
    hc_pre_eps: float = 1e-6,
    hc_sinkhorn_eps: float = 1e-6,
    hc_post_mult_value: float = 1.0,
    sinkhorn_repeat: int = 20,
) -> None: ...


def mhc_pre_fake(
    residual: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float = 1e-6,
    hc_pre_eps: float = 1e-6,
    hc_sinkhorn_eps: float = 1e-6,
    hc_post_mult_value: float = 1.0,
    sinkhorn_repeat: int = 20,  # if 0, only do pre for hc_head
    use_tf32: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    m = residual.size(0)
    hc_mult = residual.size(1)
    hidden_size = residual.size(2)
    device = residual.device
    post_mix = torch.empty(m, hc_mult, 1, dtype=dtypes.fp32, device=device)
    comb_mix = torch.empty(m, hc_mult, hc_mult, dtype=dtypes.fp32, device=device)
    layer_input = torch.empty(m, hidden_size, dtype=dtypes.bf16, device=device)
    return post_mix, comb_mix, layer_input


@torch_compile_guard(gen_fake=mhc_pre_fake)
def mhc_pre(
    residual: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float = 1e-6,
    hc_pre_eps: float = 1e-6,
    hc_sinkhorn_eps: float = 1e-6,
    hc_post_mult_value: float = 1.0,
    sinkhorn_repeat: int = 20,  # if 0, only do pre for hc_head
    use_tf32: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    m = residual.size(0)
    hc_mult = residual.size(1)
    hidden_size = residual.size(2)
    hc_mult3 = fn.size(0)
    assert hc_mult3 == hc_mult * 2 + hc_mult * hc_mult or (
        hc_mult3 == hc_mult and sinkhorn_repeat == 0
    )
    hc_hidden_size = hc_mult * hidden_size
    gfx = get_gfx()
    stage1_variant = os.environ.get("AITER_MHC_PRE_STAGE1", "auto").strip().lower()
    use_stage1_m128_auto = (
        sinkhorn_repeat > 0
        and hc_mult3 == hc_mult * (2 + hc_mult)
        and gfx != "gfx936"
        and not (hidden_size in (1280, 2560) and m <= 512)
    )
    if stage1_variant in ("", "auto"):
        use_stage1_m128 = use_stage1_m128_auto
    elif stage1_variant in ("aiter", "legacy"):
        use_stage1_m128 = False
    elif stage1_variant in ("m128", "tlstyle"):
        use_stage1_m128 = True
    else:
        raise ValueError("AITER_MHC_PRE_STAGE1 must be 'auto' or 'm128' ('tlstyle' is accepted as an alias)")

    env_kernel = os.environ.get("AITER_MHC_PRE_KERNEL", "auto").strip().lower()
    use_tlstyle_auto = (
        sinkhorn_repeat > 0
        and hc_mult3 == hc_mult * (2 + hc_mult)
        and m > 128
        and not (hidden_size in (1280, 2560) and m <= 512)
    )
    if env_kernel in ("aiter", "legacy"):
        use_tlstyle = False
    elif env_kernel == "tlstyle":
        use_tlstyle = True
    elif env_kernel in ("", "auto"):
        use_tlstyle = use_tlstyle_auto
    else:
        use_tlstyle = use_tlstyle_auto

    prefetch_stages = 2
    tile_m = 128 if use_stage1_m128 else 16 * 4
    # tile_k → 估算 tg_per_cu (target groups per CU, 受 LDS/VGPR 占用约束):
    #   tile_k=64:  tile_n*64*4*2 = 16KB/block  → 4 blocks/CU
    #   tile_k=128: tile_n*128*4*2 = 32KB/block → 2 blocks/CU
    tile_k_tg_dict = {128: 2} if use_stage1_m128 else {128: 2, 64: 4}
    num_cu = get_cu_num()
    selected_splitk = 1
    selected_tile_k = 128 if use_stage1_m128 else 64
    num_tg_m = (m + tile_m - 1) // tile_m
    # Data-driven split-k window:
    # - For small/medium M (num_tg_m < num_cu), keep broad search [1, 32].
    # - Once M-side TGs already cover all CUs (num_tg_m >= num_cu), prefer split-k=2.
    #   This avoids the large regression observed with split-k=1 on large batches.
    if num_tg_m >= num_cu:
        min_splitk = 2
        max_splitk = 2
    else:
        min_splitk = 1
        max_splitk = 32
    selected_score = num_tg_m / (num_cu * tile_k_tg_dict[selected_tile_k])
    selected_score = selected_score / math.ceil(selected_score)
    for tile_k, tg_per_cu in tile_k_tg_dict.items():
        if (hc_hidden_size % tile_k) != 0:
            continue
        meanwhile_tg = num_cu * tg_per_cu
        for splitk in range(min_splitk, max_splitk + 1):
            if hc_hidden_size % (splitk * tile_k) != 0 or (hc_hidden_size // splitk) < (
                tile_k * prefetch_stages
            ):
                continue
            num_tg = num_tg_m * splitk
            score = num_tg / meanwhile_tg
            score = score / math.ceil(score)
            if selected_score < score:
                selected_splitk = splitk
                selected_tile_k = tile_k
                selected_score = score
            # print(f"{selected_score=} {selected_splitk=} {selected_tile_k=} {score=} {splitk=} {tile_k=}")
            if num_tg > meanwhile_tg * 4:
                break

    # TileLang-style M128 stage1 still needs split-k parallelism when M-side
    # CTAs under-fill DCU. Once M-side CTAs already cover CUs, keep split_k low
    # to avoid excessive partial writes and stage2 reduction work.
    if use_stage1_m128 and hc_hidden_size in (4 * 4096, 4 * 7168):
        if num_tg_m >= num_cu:
            candidate_splitk = 2
        elif m >= 2048:
            candidate_splitk = 8
        else:
            candidate_splitk = 32
        if (
            hc_hidden_size % (candidate_splitk * selected_tile_k) == 0
            and (hc_hidden_size // candidate_splitk) >= selected_tile_k * prefetch_stages
        ):
            selected_splitk = candidate_splitk

    # Work-bound regime override:
    #   When num_tg_m >= num_cu the splitk window is already forced to {2}, and both
    #   (tile_k=64, splitk=2) and (tile_k=128, splitk=2) can land on score==1.0. The
    #   strict `<` update in the loop above lets whichever is iterated first win.
    #   Empirically on DCU gfx936/938 tile_k=64 is meaningfully faster in this regime
    #   because it halves per-block LDS occupancy (tile_n*64*4*2 vs tile_n*128*4*2),
    #   unlocking ~2x concurrent blocks per CU. Measured stage1 wins (auto vs forced
    #   tile_k=64) up to ~40% at m=8192,hidden=7168 and consistent ~10% at m=8192
    #   across hidden_size; large-m/large-hidden cases where auto already picks
    #   tile_k=64 are unchanged.
    if not use_stage1_m128 and num_tg_m >= num_cu and selected_tile_k == 128:
        candidate_tile_k = 64
        candidate_splitk = 2
        if (
            hc_hidden_size % (candidate_splitk * candidate_tile_k) == 0
            and (hc_hidden_size // candidate_splitk)
            >= candidate_tile_k * prefetch_stages
        ):
            selected_tile_k = candidate_tile_k
            selected_splitk = candidate_splitk

    # Small/medium DeepSeek MHC stage1 override:
    # sweep data shows tile_k=64, splitk=32 wins for m<=1024 on hidden=4096/7168.
    # For m=2048 it only wins on hidden=7168; hidden=4096 regresses from extra split-k work.
    candidate_tile_k = 64
    candidate_splitk = 32
    if (
        not use_stage1_m128
        and hc_hidden_size in (4 * 4096, 4 * 7168)
        and (m <= 1024 or (m == 2048 and hc_hidden_size == 4 * 7168))
        and hc_hidden_size % (candidate_splitk * candidate_tile_k) == 0
        and (hc_hidden_size // candidate_splitk) >= candidate_tile_k * prefetch_stages
    ):
        selected_tile_k = candidate_tile_k
        selected_splitk = candidate_splitk

    # Optional manual overrides for stage1 launch search:
    #   AITER_MHC_PRE_TILE_K=64|128
    #   AITER_MHC_PRE_SPLITK=<positive int>
    env_tile_k = os.environ.get("AITER_MHC_PRE_TILE_K", "").strip()
    if env_tile_k:
        forced_tile_k = int(env_tile_k)
        if forced_tile_k not in tile_k_tg_dict:
            msg = "AITER_MHC_PRE_TILE_K must be 128 when AITER_MHC_PRE_STAGE1=m128"
            if not use_stage1_m128:
                msg = "AITER_MHC_PRE_TILE_K must be 64 or 128"
            raise ValueError(msg)
        if (hc_hidden_size % forced_tile_k) != 0:
            raise ValueError(
                f"AITER_MHC_PRE_TILE_K={forced_tile_k} is incompatible with hc_hidden_size={hc_hidden_size}"
            )
        selected_tile_k = forced_tile_k

    env_splitk = os.environ.get("AITER_MHC_PRE_SPLITK", "").strip()
    if env_splitk:
        forced_splitk = int(env_splitk)
        if forced_splitk < 1:
            raise ValueError("AITER_MHC_PRE_SPLITK must be >= 1")
        if hc_hidden_size % (forced_splitk * selected_tile_k) != 0:
            raise ValueError(
                "AITER_MHC_PRE_SPLITK is incompatible with selected tile_k/hc_hidden_size"
            )
        if (hc_hidden_size // forced_splitk) < (selected_tile_k * prefetch_stages):
            raise ValueError(
                "AITER_MHC_PRE_SPLITK violates prefetch stage constraint for selected tile_k"
            )
        selected_splitk = forced_splitk

    device = residual.device
    out_pad = torch.empty(
        selected_splitk, m, (hc_mult3 + 31) // 32 * 32, dtype=dtypes.fp32, device=device
    )
    out = out_pad[:, :, :hc_mult3]
    sqrsum = torch.empty(selected_splitk, m, dtype=dtypes.fp32, device=device)
    if use_stage1_m128:
        mhc_pre_gemm_sqrsum_stage1_m128(out, sqrsum, residual, fn, use_tf32)
    else:
        stage1_fn = _round_to_tf32_like_tilekernels(fn) if use_tf32 else fn
        mhc_pre_gemm_sqrsum(out, sqrsum, residual, stage1_fn, selected_tile_k, False)
    # Optional path: reduce split-k outputs before big_fuse and run stage2 with n_splits=1.
    # Keep stage2 input layout compatible with kernel assumptions (3D + padded stride),
    # instead of passing compact 2D tensors from direct sum().
    # Enable explicitly via AITER_MHC_PRE_REDUCE_SPLITK=1|true|yes|on.
    # Current data shows the extra kernel cost outweighs the stage2 reduction win.
    use_reduce_splitk = selected_splitk > 1 and _truthy_env("AITER_MHC_PRE_REDUCE_SPLITK")
    if use_reduce_splitk:
        out_red_pad = torch.empty(
            1, m, (hc_mult3 + 31) // 32 * 32, dtype=dtypes.fp32, device=device
        )
        out_red = out_red_pad[:, :, :hc_mult3]
        sqrsum_red = torch.empty(1, m, dtype=dtypes.fp32, device=device)
        mhc_pre_reduce_splitk(out_red, sqrsum_red, out, sqrsum)
        out = out_red
        sqrsum = sqrsum_red

    post_mix = torch.empty(m, hc_mult, 1, dtype=dtypes.fp32, device=device)
    comb_mix = torch.empty(m, hc_mult, hc_mult, dtype=dtypes.fp32, device=device)
    layer_input = torch.empty(m, hidden_size, dtype=dtypes.bf16, device=device)
    big_fuse = mhc_pre_big_fuse_tlstyle if use_tlstyle else mhc_pre_big_fuse
    big_fuse(
        post_mix,
        comb_mix,
        layer_input,
        out,
        sqrsum,
        hc_scale,
        hc_base,
        residual,
        rms_eps,
        hc_pre_eps,
        hc_sinkhorn_eps,
        hc_post_mult_value,
        sinkhorn_repeat,
    )

    return post_mix, comb_mix, layer_input


@compile_ops("module_mhc")
def mhc_post(
    out: Tensor,
    x: Tensor,
    residual: Tensor,
    post_layer_mix: Tensor,
    comb_res_mix: Tensor,
) -> None: ...
