# SPDX-License-Identifier: MIT
# Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.

import torch
from torch import Tensor
from typing import Optional,List
from ..jit.core import (
    compile_ops,
)
from .enum import ActivationType, Enum, QuantType
import os
import json

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
_SILU_CFG_DIR = os.path.join(SCRIPT_DIR, "../moe_c_configs/silu_configs")
_SILU_SUMMARY = os.path.join(_SILU_CFG_DIR, "silu_config_summary.json")
_SILU_CASES_CACHE = None
_SILU_INDEX_BY_N_CACHE = None


def _load_silu_summary():
    global _SILU_CASES_CACHE, _SILU_INDEX_BY_N_CACHE
    if _SILU_CASES_CACHE is not None and _SILU_INDEX_BY_N_CACHE is not None:
        return _SILU_CASES_CACHE, _SILU_INDEX_BY_N_CACHE
    if not os.path.exists(_SILU_SUMMARY):
        _SILU_CASES_CACHE = {}
        _SILU_INDEX_BY_N_CACHE = {}
        return _SILU_CASES_CACHE, _SILU_INDEX_BY_N_CACHE
    with open(_SILU_SUMMARY, "r", encoding="utf-8") as f:
        data = json.load(f)
    _SILU_CASES_CACHE = data.get("cases", {})
    _SILU_INDEX_BY_N_CACHE = data.get("index_by_n", {})
    return _SILU_CASES_CACHE, _SILU_INDEX_BY_N_CACHE


def load_silu_tune_config(M: int, N: int):
    # 1) 只读 summary：优先精确命中 key
    N = int(N)
    cases, index_by_n = _load_silu_summary()
    key = f"M={M},N={N}"
    

    if key in cases:
        return cases[key]["rows_per_block"], cases[key]["vec_size"]

    # 2) 同 N 下按 |M - M_i| 找最接近配置（优先使用预构建索引）
    n_key = str(N)
    if n_key in index_by_n and index_by_n[n_key]:
        entries = index_by_n[n_key]  # sorted by M
        m_list = [int(e["M"]) for e in entries]

        # Manual lower_bound (avoid importing bisect).
        left = 0
        right = len(m_list)
        while left < right:
            mid = (left + right) // 2
            if m_list[mid] < M:
                left = mid + 1
            else:
                right = mid
        pos = left

        candidates = []
        if pos < len(entries):
            candidates.append(entries[pos])
        if pos > 0:
            candidates.append(entries[pos - 1])
        if candidates:
            best = min(candidates, key=lambda e: abs(int(e["M"]) - M))
            return best["rows_per_block"], best["vec_size"]

    # Backward-compatible slow path when old summary has no index_by_n.
    nearest = None
    for _, v in cases.items():
        if int(v.get("N", -1)) != N:
            continue
        km = int(v.get("M", -1))
        if km < 0:
            continue
        dist = abs(km - M)
        if nearest is None or dist < nearest[0]:
            nearest = (dist, v)
    if nearest is not None:
        return nearest[1]["rows_per_block"], nearest[1]["vec_size"]

    # 3) fallback 默认值
    
    return 1, 2


@compile_ops("module_moe_c_kernel")
def moe_c_moe_gemm_marlin_w8a8( 
    input: torch.Tensor,
    b_qweight : torch.Tensor,
    output : torch.Tensor,
    a_scale: torch.Tensor,
    b_scale : torch.Tensor,
    topk_weights : Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids : torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k : int,
    mode :int,
    delta: int,
    size_m: int)-> torch.Tensor: 
    """
    ---------------------------------------------------------------
    # MoE 场景下 8bit 量化的 GEMM 计算（Marlin 优化版）   

    ## 关键前置条件
    必须配合对应的权重 Shuffle 函数使用，否则会导致计算结果完全错误：
    - GEMM1 场景：使用 ops.marlin_weights 处理权重
    - GEMM2 场景：使用 ops.marlin_weights_ours 处理权重




    
    ---------------------------------------------------------------
    """
    
    pass


@compile_ops("module_moe_c_kernel")
def moe_c_moe_gemm_marlin_w8a8_tensorwise( 
    input: torch.Tensor,
    b_qweight : torch.Tensor,
    output : torch.Tensor,
    a_scale: torch.Tensor,
    b_scale : torch.Tensor,
    topk_weights : Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids : torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k : int,
    mode :int,
    delta: int,
    size_m: int)-> torch.Tensor: 
    """
    Marlin W8A8 MoE GEMM with tensorwise weight scales.

    b_scale must contain one scale per expert and use shape (E, 1, 1).
    """
    
    pass


@compile_ops("module_moe_c_kernel")
def moe_c_moe_gemm_marlin_w4a8( 
    input: torch.Tensor,
    b_qweight : torch.Tensor,
    output : torch.Tensor,
    a_scale: torch.Tensor,
    b_scale : torch.Tensor,
    topk_weights : Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids : torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k : int,
    mode :int,
    delta: int,
    size_m: int)-> torch.Tensor: 
    """
    ---------------------------------------------------------------
    # MoE 场景下 8bit 量化的 GEMM 计算（Marlin 优化版）   

    ## 关键前置条件
    必须配合对应的权重 Shuffle 函数使用，否则会导致计算结果完全错误：
    - GEMM1 场景：使用 ops.marlin_weights 处理权重
    - GEMM2 场景：使用 ops.marlin_weights_ours 处理权重




    
    ---------------------------------------------------------------
    """
    
    pass
      
    

@compile_ops("module_moe_c_kernel")
def moe_c_moe_gemm_marlin_w8a8_fp8( 
    input: torch.Tensor,
    b_qweight : torch.Tensor,
    output : torch.Tensor,
    a_scale: torch.Tensor,
    b_scale : torch.Tensor,
    topk_weights : Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids : torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k : int,
    mode :int,
    delta: int,
    size_m: int)-> torch.Tensor: 
    """
    ---------------------------------------------------------------
    # MoE 场景下 8bit 量化的 GEMM 计算（Marlin 优化版）   

    ## 关键前置条件
    必须配合对应的权重 Shuffle 函数使用，否则会导致计算结果完全错误：
    - GEMM1 场景：使用 ops.marlin_weights 处理权重
    - GEMM2 场景：使用 ops.marlin_weights_ours 处理权重




    
    ---------------------------------------------------------------
    """
    
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_gemm_marlin_w8a8_fp8_tensorwise( 
    input: torch.Tensor,
    b_qweight : torch.Tensor,
    output : torch.Tensor,
    a_scale: torch.Tensor,
    b_scale : torch.Tensor,
    topk_weights : Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids : torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k : int,
    mode :int,
    delta: int,
    size_m: int)-> torch.Tensor: 
    """
    Marlin FP8 W8A8 MoE GEMM with tensorwise weight scales.

    b_scale must contain one scale per expert and use shape (E, 1, 1).
    """
    
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_gemm_marlin_w4a16( 
    input: torch.Tensor,
    b_qweight : torch.Tensor,
    output : torch.Tensor,
    b_scale: torch.Tensor,
    b_zeros : torch.Tensor,
    topk_weights : Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids : torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k : int,
    mode :int,
    delta: int)-> torch.Tensor: 
    """
    ---------------------------------------------------------------
    # MoE 场景下 4bit 量化的 GEMM 计算（Marlin 优化版）   

    ## 关键前置条件
    必须配合对应的权重 Shuffle 函数使用，否则会导致计算结果完全错误：

    
    ---------------------------------------------------------------
    """
    
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_gemm_marlin_w8a16( 
    input: torch.Tensor,
    b_qweight : torch.Tensor,
    output : torch.Tensor,
    b_scale: torch.Tensor,
    topk_weights : Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids : torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k : int,
    mode :int,
    delta: int)-> torch.Tensor: 
    """
    ---------------------------------------------------------------
    # MoE 场景下 4bit 量化的 GEMM 计算（Marlin 优化版）   

    ## 关键前置条件
    必须配合对应的权重 Shuffle 函数使用，否则会导致计算结果完全错误：

    
    ---------------------------------------------------------------
    """
    
    pass


@compile_ops("module_moe_c_kernel")
def moe_c_moe_w8a8_gemm_block_wise(
    input: torch.Tensor,
    a_scales: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],   
    topk_weights: Optional[torch.Tensor],   
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    group_size_n: int,  
    group_size_k: int,  
    top_k: int,   
    BLOCK_SIZE_m: int,  
    BLOCK_SIZE_n: int,  
    BLOCK_SIZE_k: int,  
    kloops: int,   
    nloops: int,   
    bit: int   
) -> torch.Tensor:  
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_w8a8_gemm_block_wise_kernel2(
    input: torch.Tensor,
    a_scales: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    group_size_n: int,
    group_size_k: int,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    kloops: int,
    nloops: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_w8a8_gemm_block_wise_fp8(
    input: torch.Tensor,
    a_scales: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    group_size_n: int,
    group_size_k: int,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    kloops: int,
    nloops: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_w8a8_gemm_block_wise_kernel2_fp8(
    input: torch.Tensor,
    a_scales: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    group_size_n: int,
    group_size_k: int,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    kloops: int,
    nloops: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_w8a16_gemm_awq(
    input: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_w8a16_gemm_block_wise(
    input: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    group_size_n: int,
    group_size_k: int,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_wna16_gemm_base(
    input: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k: int,
    BLOCK_SIZE_M: int,
    BLOCK_SIZE_N: int,
    BLOCK_SIZE_K: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_wna16_gemm(
    input: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    kloops: int,
    nloops: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_wna16_gemm_2(
    input: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    kloops: int,
    nloops: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_topk_softmax(
    topk_weights: torch.Tensor,  # 移除 C++ 引用 &
    topk_indices: torch.Tensor,  # 移除 C++ 引用 &
    token_expert_indices: torch.Tensor,  # 移除 C++ 引用 &
    gating_output: torch.Tensor  # 移除 C++ 引用 &
) -> None:  # 替代 -> None (C++ 中的 void)
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_silu_and_mul( out : torch.Tensor,   
                  input : torch.Tensor,
                  rows_per_block: int = 1, 
                  vec_size: int = 2) -> None:
    pass





@compile_ops("module_moe_c_kernel")
def moe_c_moe_sum(
    input: torch.Tensor,  # 移除 C++ 引用 &
    output: torch.Tensor,  # 移除 C++ 引用 &
    topk_ids: torch.Tensor
) -> None:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_sum_opt_v2(input: torch.Tensor,output: torch.Tensor,
                   routed_scaling_factor: float = 1.0) -> torch.Tensor:
    pass



@compile_ops("module_moe_c_kernel")
def moe_c_moe_align_block_size(
    topk_ids: torch.Tensor,
    num_experts: int,
    block_size: int,
    sorted_token_ids: torch.Tensor,
    experts_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor
) -> None:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_sgl_moe_align_block_size(
    topk_ids: torch.Tensor,
    num_experts: int,
    block_size: int,
    sorted_token_ids: torch.Tensor,
    experts_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor
) -> None:
    pass



