import functools
import os
import re

import torch

_ENABLED = os.getenv("MAGI_TIRX_DOWN", "0") == "1"
_VALIDATE = os.getenv("MAGI_TIRX_DOWN_VALIDATE", "0") == "1"
_checked = set()
_reported = False

KP = 13696
N = 5120


@functools.lru_cache(maxsize=2)
def executable(mp):
    import tirx_kernels.gemm.fp16_bf16_gemm as gemm
    from tirx_kernels.runner import compile_kernel

    cfg = dict(gemm.GEMM_CONFIGS[1024])
    cfg.update(
        cta_n=64,
        cta_k=128,
        pipe_depth=5,
        wb_pipe_depth=2,
        l2_group_size=2,
        overlap_epilogue=True,
    )

    original_cfg = gemm._cfg_for
    try:
        gemm._cfg_for = lambda n: dict(cfg)
        kernel = gemm.make_kernel("bf16", mp, N, KP)
        result = compile_kernel(kernel)
    finally:
        gemm._cfg_for = original_cfg

    print(f"TIRX_DOWN_COMPILED M={mp} N={N} K={KP}", flush=True)
    return result


def supported(module, z, weight):
    return (
        _ENABLED
        and not torch.is_grad_enabled()
        and not module.training
        and z.is_cuda
        and weight.device == z.device
        and z.dtype == torch.bfloat16
        and weight.dtype == torch.bfloat16
        and z.ndim == 3
        and weight.ndim == 3
        and z.shape[0] == 3
        and 0 < z.shape[1] <= 4096
        and z.shape[2] == 13652
        and tuple(weight.shape) == (3, 5120, 13652)
        and z.is_contiguous()
        and weight.is_contiguous()
        and not z.requires_grad
        and not weight.requires_grad
    )


def down(module, z, weight):
    global _reported

    m = z.shape[1]
    mp = 2048 if m <= 2048 else 4096

    # 保存原权重对象，避免地址复用导致误用缓存。
    # 权重被原地修改后，通过 version 重新建立缓存。
    # inference tensor没有版本计数；固定权重推理按对象身份缓存。
    version = None if torch.is_inference(weight) else weight._version
    cached = getattr(module, "_tirx_down_weight_cache", None)
    if (
        cached is None
        or cached[0] is not weight
        or cached[1] != version
    ):
        padded_weight = torch.zeros(
            3, N, KP, dtype=weight.dtype, device=weight.device
        )
        padded_weight[:, :, :13652].copy_(weight)
        module._tirx_down_weight_cache = (
            weight, version, padded_weight
        )
    else:
        padded_weight = cached[2]

    # 全部清零，确保 K 尾部和 M 尾部没有未初始化值。
    padded_input = torch.zeros(
        3, mp, KP, dtype=z.dtype, device=z.device
    )
    padded_input[:, :m, :13652].copy_(z)

    padded_output = torch.empty(
        3, mp, N, dtype=z.dtype, device=z.device
    )
    ex = executable(mp)
    for e in range(3):
        ex(padded_input[e], padded_weight[e], padded_output[e])

    output = padded_output[:, :m, :].contiguous()

    if not _reported:
        print(
            f"TIRX_DOWN_ACTIVE input={tuple(z.shape)} "
            f"padded_M={mp} padded_K={KP}",
            flush=True,
        )
        _reported = True

    # 检查前三种真实形状；这不是全模型质量验证。
    key = tuple(z.shape)
    if _VALIDATE and key not in _checked and len(_checked) < 3:
        expected = torch.bmm(z, weight.transpose(-2, -1))
        finite = bool(torch.isfinite(output).all().item())

        # 分块统计，控制验证时的临时显存。
        a = output.reshape(-1)
        b = expected.reshape(-1)
        max_abs = 0.0
        error_sq = 0.0
        reference_sq = 0.0
        for start in range(0, a.numel(), 1 << 20):
            af = a[start:start + (1 << 20)].float()
            bf = b[start:start + (1 << 20)].float()
            delta = af - bf
            max_abs = max(max_abs, delta.abs().max().item())
            error_sq += delta.square().sum(dtype=torch.float64).item()
            reference_sq += bf.square().sum(dtype=torch.float64).item()

        relative_l2 = (error_sq / max(reference_sq, 1e-30)) ** 0.5
        print(
            f"TIRX_DOWN_CHECK input={tuple(z.shape)} "
            f"finite={finite} max_abs={max_abs:.8g} "
            f"relative_l2={relative_l2:.8g}",
            flush=True,
        )

        # 仅用于拦截明显异常，不能代替质量评价。
        if not finite or relative_l2 > 1e-3:
            raise RuntimeError("TIRx down真实输入检查未通过")

        _checked.add(key)

    return output


def transform_source(source):
    pattern = re.compile(
        r"^(?P<indent>[ \t]*)"
        r"y\s*=\s*torch\.bmm\("
        r"z,\s*w_dn\.transpose\(-2,\s*-1\)"
        r"\)\.to\(dtype\)[ \t]*$",
        re.MULTILINE,
    )
    matches = list(pattern.finditer(source))
    if len(matches) != 1:
        raise RuntimeError(
            f"down BMM入口匹配到{len(matches)}处，未安装补丁"
        )

    match = matches[0]
    indent = match.group("indent")
    replacement = (
        indent + "if _tirx_down_supported(self, z, w_dn):\n"
        + indent + "    y = _tirx_down(self, z, w_dn)\n"
        + indent + "else:\n"
        + indent
        + "    y = torch.bmm(z, w_dn.transpose(-2, -1)).to(dtype)"
    )
    return source[:match.start()] + replacement + source[match.end():]


@torch.inference_mode()
def prewarm_model(model):
    from inference.model.dit.dit_module import SparseMoEMLP

    modules = [
        m for m in model.modules()
        if isinstance(m, SparseMoEMLP)
    ]
    if not modules:
        raise RuntimeError("down预热未找到SparseMoEMLP")

    print(
        f"TIRX_PREWARM_DOWN: 开始 modules={len(modules)} "
        f"enabled={_ENABLED}",
        flush=True,
    )

    first = True
    for index, module in enumerate(modules):
        if module.training:
            raise RuntimeError("down预热要求eval模式")

        experts = module.experts
        device = experts[0].down_proj.weight.device
        if device.type != "cuda":
            raise RuntimeError("down预热要求模型已在GPU，当前不支持CPU offload")

        if (
            len(experts) != 3
            or any(
                tuple(e.down_proj.weight.shape) != (5120, 13652)
                for e in experts
            )
        ):
            raise RuntimeError("down预热只适用于当前三专家、全宽13652模型")

        # 与原始 _cache_key("bmm", None) 完全一致。
        key = ("bmm", torch.bfloat16, device.type, device.index, -1)
        cache = module._grouped_weight_cache

        if key not in cache:
            up = torch.stack(
                [e.up_gate_proj.weight.to(torch.bfloat16) for e in experts],
                dim=0,
            )
            dn = torch.stack(
                [e.down_proj.weight.to(torch.bfloat16) for e in experts],
                dim=0,
            )
            cache[key] = (up, dn)

        if _ENABLED:
            weight = cache[key][1]
            version = (
                None if torch.is_inference(weight) else weight._version
            )
            padded = torch.zeros(
                3, N, KP, dtype=torch.bfloat16, device=device
            )
            padded[:, :, :13652].copy_(weight)
            module._tirx_down_weight_cache = (weight, version, padded)

            # 两档内核各启动一次；不调用down()，避免消耗真实输入检查名额。
            if first:
                for mp in (2048, 4096):
                    x = torch.zeros(
                        mp, KP, dtype=torch.bfloat16, device=device
                    )
                    y = torch.empty(
                        mp, N, dtype=torch.bfloat16, device=device
                    )
                    ex = executable(mp)
                    for _ in range(3):
                        ex(x, padded[0], y)
                    torch.cuda.synchronize(device)
                    del x, y
                first = False

        print(
            f"TIRX_PREWARM_DOWN: 权重准备 {index + 1}/{len(modules)}",
            flush=True,
        )

    torch.cuda.synchronize()
    print("TIRX_PREWARM_DOWN: 完成", flush=True)