"""Kernel classification logic for trace comparison.

Classifies GPU kernels into operation categories (attention, GEMM, normalization, etc.)
based on kernel name patterns and platform-specific conventions.
"""

from enum import Enum


class Op(Enum):
    """Kernel operation categories."""

    ATTENTION = "Attention"
    KV_CACHE = "KV Cache"
    MOE_ROUTING = "MoE Routing"
    MOE_GEMM = "MoE GEMM"
    MOE_GEMM_SWIGLU = "MoE GEMM+SwiGLU"
    MOE_FINALIZE = "MoE Finalize"
    DENSE_GEMM = "Dense GEMM"
    RMSNORM = "RMSNorm"
    RMSNORM_GEMM = "RMSNorm+GEMM"
    TRITON_FUSED = "Triton Fused"
    ELEMENTWISE = "Elementwise"
    SORTING = "Sorting"
    REDUCE = "Reduce"
    COPY_MEMORY = "Copy/Memory"
    COMMUNICATION = "Communication"
    ALL_REDUCE = "All-Reduce"
    ALL_GATHER = "All-Gather"
    OTHER = "Other"


def classify(name: str, platform: str) -> tuple[Op, str]:
    """Classify kernel by operation type.

    Platform-specific classification ensures accurate semantic mapping across AMD and NVIDIA.
    Avoid over-broad patterns that could cause false fusion detection.

    Args:
        name: Kernel name from trace
        platform: 'AMD' or 'NVIDIA'

    Returns:
        Tuple of (operation type, pattern name)
    """
    nl = name.lower()

    if "attention" in nl or "fmha" in nl:
        if platform == "AMD":
            if "2d" in nl:
                return Op.ATTENTION, "kernel_unified_attention_2d"
            if "3d" in nl:
                return Op.ATTENTION, "kernel_unified_attention_3d"
        else:
            if "fmhasm100a" in nl or "context" in nl:  # 'a'=prefill/context, 'f'=decode/forgen
                return Op.ATTENTION, "fmhaSm100a*_Context"
            if "fmhasm100f" in nl or "forgen" in nl:
                return Op.ATTENTION, "fmhaSm100f*_ForGen"
        return Op.ATTENTION, name[:40]

    if "reshape_and_cache" in nl:
        return Op.KV_CACHE, "reshape_and_cache_*"
    if "create_flashinfer_kv_indices" in nl or "flashinfer" in nl:
        return Op.KV_CACHE, "flashinfer_kv_*"

    if "_matmul_ogs_" in nl:
        if "swiglu" in nl:
            return Op.MOE_GEMM_SWIGLU, "_matmul_ogs_*_swiglu"
        return Op.MOE_GEMM, "_matmul_ogs_*"

    if name.startswith("bmm_") and "dynbatch" in nl:
        if "swiglu" in nl:
            return Op.MOE_GEMM_SWIGLU, "bmm_*_swiGlu_dynBatch"
        return Op.MOE_GEMM, "bmm_*_dynBatch"

    if "topk" in nl or "routing" in nl or "bitmatrix" in nl or "moe_forward" in nl or "_combined_routing" in nl:
        return Op.MOE_ROUTING, "moe_routing_*"
    if "finalize" in nl or ("scatter" in nl and "moe" in nl):
        return Op.MOE_FINALIZE, "moe_finalize_*"

    # AMD kernel names may include misleading hints (e.g. "rocm_unquantized_gemm_rsqrt" in name);
    # graph matching confirms these are standalone RMSNorm, not RMSNorm+GEMM fusions.
    if "triton" in nl and ("rsqrt" in nl or ("mean" in nl and "mul" in nl and "pow" in nl)):
        return Op.RMSNORM, "triton_*_rsqrt"
    if "rmsnorm2dfwd" in nl:
        return Op.RMSNORM, "ck_tile::Rmsnorm2dFwd"

    if name.startswith("Cijk_") or name.startswith("Custom_Cijk_"):
        return Op.DENSE_GEMM, "Cijk_* (Tensile)"
    if name.startswith("nvjet_") or "cublaslt" in nl:
        return Op.DENSE_GEMM, "nvjet_* (cuBLASLt)"
    if "wvsplitk" in nl or name.startswith("void wvSplitK"):
        return Op.DENSE_GEMM, "wvSplitK_* (hipBLASLt)"

    if "triton_poi" in nl or "triton_red" in nl or "triton_per" in nl:
        if "silu" in nl or "swiglu" in nl:
            return Op.TRITON_FUSED, "triton_*_silu"
        return Op.TRITON_FUSED, "triton_*"

    if "act_and_mul_kernel" in nl or "sgl_hip::activation" in nl:
        return Op.ELEMENTWISE, "sgl_activation_*"

    if "compute_position_kernel" in nl:
        return Op.ELEMENTWISE, "compute_position_kernel"
    if "rotary_embedding" in nl:
        return Op.ELEMENTWISE, "rotary_embedding_*"

    if "at::native::" in name:
        return Op.ELEMENTWISE, "at::native::*"

    if "sort" in nl or "radixsort" in nl or "merge" in nl:
        if platform == "AMD":
            return Op.SORTING, "rocprim::sort/merge_*"
        else:
            return Op.SORTING, "cub::DeviceRadixSort*"

    if "scan_kernel" in nl or ("rocprim" in nl and "scan" in nl):
        return Op.REDUCE, "rocprim::scan_*"
    if "devicescan" in nl or ("cub" in nl and "scan" in nl):
        return Op.REDUCE, "cub::DeviceScan*"

    # Communication operations (multi-GPU/tensor parallel)
    # NCCL all-gather operations
    if "nccl" in nl and "allgather" in nl:
        return Op.ALL_GATHER, "ncclDevKernel_AllGather_*"
    if "_all_gather_base" in nl or "all_gather" in nl:
        return Op.ALL_GATHER, "nccl:_all_gather_base"

    # All-reduce operations (custom vLLM kernels and NCCL)
    if "cross_device_reduce" in nl:
        return Op.ALL_REDUCE, "vllm::cross_device_reduce_*"
    if "all_reduce" in nl:
        return Op.ALL_REDUCE, "vllm::all_reduce"
    if "two_shot_all_reduce" in nl:
        return Op.ALL_REDUCE, "symm_mem::two_shot_all_reduce_*"

    # Generic NCCL operations
    if "nccldevkernel" in nl or "nccl::" in nl:
        return Op.COMMUNICATION, "ncclDevKernel_*"

    # Reduce operations
    if "reduce" in nl and ("reduce_segments" in nl or "devicereduce" in nl or "devicescan" in nl):
        if platform == "AMD":
            return Op.REDUCE, "reduce_segments"
        else:
            return Op.REDUCE, "cub::DeviceReduce*"

    if "copy" in nl or "memcpy" in nl or "_copy_page_indices" in nl:
        return Op.COPY_MEMORY, "copy_*"

    if "rocprim::" in name or "cub::" in name:
        return Op.OTHER, "rocprim/cub_*"

    return Op.OTHER, name[:40]


def classify_kernel(name: str) -> str:
    """Simplified kernel classification for fusion analysis.

    Args:
        name: Kernel name from trace

    Returns:
        Simple category name consistent across platforms
    """
    nl = name.lower()

    if any(x in nl for x in ["cijk_", "nvjet", "wvsplitk", "cublas", "hipblas", "tensile"]):
        return "GEMM"

    if "attention" in nl or "fmha" in nl:
        return "Attention"

    if "reshape_and_cache" in nl:
        return "KV_Cache"

    if "triton" in nl and "rsqrt" in nl:
        return "RMSNorm"
    if "layernorm" in nl or "rmsnorm" in nl:
        return "RMSNorm"

    if "silu" in nl or "swiglu" in nl:
        return "SwiGLU"
    if "gelu" in nl:
        return "GELU"
    if "relu" in nl and "gelu" not in nl:
        return "ReLU"

    if "triton_poi" in nl:
        return "Triton_Pointwise"
    if "triton_red" in nl:
        return "Triton_Reduce"
    if "triton_per" in nl:
        return "Triton_Persistent"

    if "reduce_segments" in nl or "devicereduce" in nl:
        return "Reduce"

    if "sort" in nl or "radixsort" in nl or "merge" in nl:
        return "Sort"

    if "softmax" in nl:
        return "Softmax"

    if any(x in nl for x in ["elementwise", "unrolled_elementwise"]):
        return "Elementwise"

    if "copy" in nl or "memcpy" in nl:
        return "MemCopy"

    # Communication operations (multi-GPU)
    if "all_gather" in nl:
        return "All_Gather"
    if "all_reduce" in nl or "cross_device_reduce" in nl:
        return "All_Reduce"
    if "nccl" in nl:
        return "NCCL_Collective"

    return "Other"
