"""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.

Can optionally load patterns from kernel_registry.yaml, but falls back to hardcoded patterns
for comprehensive coverage.
"""

import fnmatch
from enum import Enum
from functools import lru_cache
from pathlib import Path
from typing import Any

try:
    import yaml
except ImportError:
    yaml = None


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

    ATTN_PREFILL = "Attention (Prefill)"
    ATTN_DECODE = "Attention (Decode)"
    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"
    SWIGLU = "SwiGLU"
    SWIGLU_GEMM = "SwiGLU+GEMM"
    EMBEDDING_RMSNORM_GEMM = "Embedding+RMSNorm+GEMM"
    SOFTMAX = "SoftMax"
    TRITON_FUSED = "Triton Fused"
    ELEMENTWISE = "Elementwise"
    SORTING = "Sorting"
    REDUCE = "Reduce"
    INDEXING = "Indexing"
    COPY_MEMORY = "Copy/Memory"
    FUSED_UNKNOWN = "Fused (Unknown)"  # Heuristically detected fusion
    OTHER = "Other"


# Keywords that indicate specific operations - used for heuristic fusion detection
FUSION_KEYWORDS: dict[str, str] = {
    # Normalization
    "rsqrt": "Norm",
    "rmsnorm": "RMSNorm",
    "layernorm": "LayerNorm",
    # GEMM
    "gemm": "GEMM",
    "matmul": "GEMM",
    "mm_": "GEMM",
    # Activations
    "silu": "SiLU",
    "swiglu": "SwiGLU",
    "gelu": "GELU",
    "relu": "ReLU",
    # Other ops
    "softmax": "Softmax",
    "attention": "Attention",
    "embedding": "Embedding",
    "reduce": "Reduce",
}


# Load kernel registry YAML if available
_KERNEL_REGISTRY: dict[str, Any] | None = None


def _load_kernel_registry() -> dict[str, Any] | None:
    """Load kernel pattern registry from YAML file."""
    global _KERNEL_REGISTRY
    
    if _KERNEL_REGISTRY is not None:
        return _KERNEL_REGISTRY
    
    if yaml is None:
        return None
    
    registry_path = Path(__file__).parent / "kernel_registry.yaml"
    if not registry_path.exists():
        return None
    
    try:
        with open(registry_path) as f:
            _KERNEL_REGISTRY = yaml.safe_load(f)
        return _KERNEL_REGISTRY
    except Exception:
        return None


def _match_registry_pattern(name: str, category: str, platform: str) -> tuple[Op, str] | None:
    """Try to match kernel name against registry patterns.
    
    Returns (Op, pattern_name) if match found, None otherwise.
    """
    registry = _load_kernel_registry()
    if not registry:
        return None
    
    nl = name.lower()
    category_data = registry.get(category, {})
    
    platform_key = "amd" if platform.lower() == "amd" else "nvidia"
    patterns = category_data.get(platform_key, [])
    both_patterns = category_data.get("both", [])
    patterns = patterns + both_patterns
    
    for pattern_entry in patterns:
        pattern = pattern_entry.get("pattern", "")
        if not pattern:
            continue
        
        if fnmatch.fnmatch(name, pattern) or fnmatch.fnmatch(nl, pattern.lower()):
            if category == "attention":
                phase = pattern_entry.get("phase")
                if phase == "prefill":
                    return Op.ATTN_PREFILL, pattern
                elif phase == "decode":
                    return Op.ATTN_DECODE, pattern
                return Op.ATTN_PREFILL, pattern
            elif category == "gemm":
                return Op.DENSE_GEMM, pattern
            elif category == "rmsnorm":
                # Detect fused operations (AMD Triton often fuses RMSNorm with GEMM)
                has_gemm = "gemm" in nl or "unquantized_gemm" in nl
                has_embedding = "embedding" in nl
                if has_gemm and has_embedding:
                    # Embedding + RMSNorm + GEMM all fused together
                    return Op.EMBEDDING_RMSNORM_GEMM, pattern
                elif has_gemm:
                    # RMSNorm + GEMM fused
                    return Op.RMSNORM_GEMM, pattern
                return Op.RMSNORM, pattern
            elif category == "moe":
                # Distinguish between MoE sub-operations
                if "swiglu" in nl:
                    return Op.MOE_GEMM_SWIGLU, pattern
                if any(x in nl for x in ["routing", "topk", "align_block", "count_and_sort", "gating"]):
                    return Op.MOE_ROUTING, pattern
                if "finalize" in nl or "scatter" in nl:
                    return Op.MOE_FINALIZE, pattern
                return Op.MOE_GEMM, pattern
            elif category == "kv_cache":
                return Op.KV_CACHE, pattern
            elif category == "softmax":
                return Op.SOFTMAX, pattern
            elif category == "reduce":
                return Op.REDUCE, pattern
            elif category == "sorting":
                return Op.SORTING, pattern
            elif category == "memory":
                return Op.COPY_MEMORY, pattern
            elif category == "indexing":
                return Op.INDEXING, pattern
            elif category == "elementwise":
                return Op.ELEMENTWISE, pattern
            elif category == "triton":
                return Op.TRITON_FUSED, pattern
            elif category == "activation":
                # Check for fused SwiGLU+GEMM (AMD Triton)
                has_gemm = "gemm" in nl or "unquantized_gemm" in nl
                has_silu = "silu" in nl or "swiglu" in nl
                if has_gemm and has_silu:
                    return Op.SWIGLU_GEMM, pattern
                elif has_silu:
                    return Op.SWIGLU, pattern
                # Other activations (GELU, etc.)
                return Op.TRITON_FUSED, pattern
    
    return None


def _detect_heuristic_fusion(name: str) -> tuple[Op, str] | None:
    """Heuristically detect potential fusions based on multiple operation keywords.
    
    This is a fallback for kernels we haven't explicitly classified.
    If a kernel name contains 2+ distinct operation keywords, it's likely fused.
    
    Returns (Op.FUSED_UNKNOWN, "Component1+Component2+...") if suspected fusion.
    The pattern name contains the fused components for display.
    """
    nl = name.lower()
    
    # Only check Triton kernels - these are most likely to be fused
    if "triton" not in nl:
        return None
    
    # Find all operation keywords present in the name
    # Use ordered list to maintain consistent ordering
    found_ops: list[str] = []
    keyword_priority = [
        # Order matters - more specific first
        ("embedding", "Embedding"),
        ("rmsnorm", "RMSNorm"),
        ("layernorm", "LayerNorm"),
        ("rsqrt", "Norm"),  # Generic norm indicator
        ("swiglu", "SwiGLU"),
        ("silu", "SiLU"),
        ("gelu", "GELU"),
        ("relu", "ReLU"),
        ("gemm", "GEMM"),
        ("matmul", "GEMM"),
        ("mm_", "GEMM"),
        ("softmax", "Softmax"),
        ("attention", "Attention"),
        ("reduce", "Reduce"),
    ]
    
    for keyword, op_name in keyword_priority:
        if keyword in nl and op_name not in found_ops:
            # Avoid duplicates like "RMSNorm" and "Norm" 
            if op_name == "Norm" and any(n in found_ops for n in ["RMSNorm", "LayerNorm"]):
                continue
            # Avoid duplicates like "SwiGLU" and "SiLU"
            if op_name == "SiLU" and "SwiGLU" in found_ops:
                continue
            found_ops.append(op_name)
    
    # If 2+ operations detected, it's likely a fusion
    if len(found_ops) >= 2:
        fused_name = "+".join(found_ops)
        # The pattern name IS the fused operation name for display
        return Op.FUSED_UNKNOWN, fused_name
    
    return None


@lru_cache(maxsize=4096)
def classify(name: str, platform: str) -> tuple[Op, str]:
    """Classify kernel by operation type.
    
    Cached because PyTorch traces have ~48 unique kernel names repeated 810k times.

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

    Returns:
        Tuple of (operation type, pattern name)
    """
    nl = name.lower()
    
    # Check registry patterns first (order matters - more specific categories first)
    registry_categories = [
        "attention",   # Attention ops (prefill/decode)
        "gemm",        # Dense GEMM
        "rmsnorm",     # RMSNorm
        "moe",         # MoE operations
        "kv_cache",    # KV Cache
        "activation",  # SwiGLU, SiLU, GELU
        "softmax",     # Softmax
        "reduce",      # Reduce/Scan
        "sorting",     # Sorting
        "memory",      # Memory/Copy
        "indexing",    # Index/Scatter-Gather
        "triton",      # Triton fused ops
        "elementwise", # Elementwise (last - most generic)
    ]
    for category in registry_categories:
        result = _match_registry_pattern(name, category, platform)
        if result:
            return result
    if "attention" in nl or "fmha" in nl:
        if platform == "AMD":
            if "2d" in nl:
                return Op.ATTN_PREFILL, "kernel_unified_attention_2d"
            if "3d" in nl:
                return Op.ATTN_DECODE, "kernel_unified_attention_3d"
        else:
            # NVIDIA uses fmhaSm100 with 'a' (prefill/context) and 'f' (decode/forgen)
            if "fmhasm100a" in nl or "context" in nl:
                return Op.ATTN_PREFILL, "fmhaSm100a*_Context"
            if "fmhasm100f" in nl or "forgen" in nl:
                return Op.ATTN_DECODE, "fmhaSm100f*_ForGen"
        return Op.ATTN_PREFILL, name[:40]

    # Flash Attention variants (vLLM)
    if "flash::flash_fwd_kernel" in name or "flash_fwd" in nl:
        # Could distinguish prefill/decode if needed, defaulting to prefill
        return Op.ATTN_PREFILL, "flash::flash_fwd_kernel"

    if "reshape_and_cache" in nl:
        return Op.KV_CACHE, "reshape_and_cache_*"

    # KV Cache variants (vLLM)
    if "concat_and_cache" in nl or "cache_mla" in nl:
        return Op.KV_CACHE, "vllm::concat_and_cache_*"

    # MoE
    # vLLM MoE kernels - these are very common in MoE models
    if "fused_moe_kernel" in nl:
        return Op.MOE_GEMM, "fused_moe_kernel"

    if "vllm::moe::" in name:
        if "moe_align_block_size" in nl:
            if "small_batch" in nl:
                return Op.MOE_ROUTING, "vllm::moe::moe_align_block_size_small_batch_*"
            return Op.MOE_ROUTING, "vllm::moe::moe_align_block_size_*"
        if "moe_sum" in nl:
            return Op.MOE_FINALIZE, "vllm::moe::moe_sum_*"

    # vLLM act_and_mul (can be mangled C++ name)
    if "vllm::act_and_mul_kernel" in name or ("act_and_mul_kernel" in nl and "vllm" in nl):
        return Op.MOE_GEMM_SWIGLU, "vllm::act_and_mul_kernel"

    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"

    # Generic MoE routing patterns (check before finalize)
    if any(x in nl for x in ["topk", "routing", "bitmatrix", "moe_forward", "_combined_routing"]):
        if "moe::dev::routing::" in name or "moe::" in name:
            return Op.MOE_ROUTING, "moe::dev::routing::*"
        return Op.MOE_ROUTING, "moe_routing_*"
    
    # MoE finalize patterns
    if "finalize" in nl or ("scatter" in nl and "moe" in nl):
        return Op.MOE_FINALIZE, "moe_finalize_*"

    # RMSNorm - match actual patterns from traces
    if "triton" in nl and ("rsqrt" in nl or ("mean" in nl and "mul" in nl and "pow" in nl)):
        if "gemm" in nl or "addmm" in nl:
            return Op.RMSNORM_GEMM, "triton_*_rmsnorm_gemm"
        return Op.RMSNORM, "triton_*_rsqrt"

    # Dense GEMM - these are the most common kernels
    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)"
    
    # CUTLASS GEMM variants
    if "cutlass" in nl and ("sgemm" in nl or "gemm" in nl or "cutlass3x" in name.lower()):
        return Op.DENSE_GEMM, "cutlass*_gemm"
    
    # GEMV (matrix-vector) operations - treat as GEMM
    if "gemv" in nl or "gemvx" in nl:
        return Op.DENSE_GEMM, "gemv*"
    
    # Generic GEMM patterns
    if "gemmsn" in nl or name.startswith("void gemmSN"):
        return Op.DENSE_GEMM, "gemmSN_*"

    # Triton fused operations - very common
    if "triton_poi" in nl or "triton_red" in nl or "triton_per" in nl:
        # Distinguish between different fusion patterns
        if "silu" in nl or "swiglu" in nl:
            return Op.TRITON_FUSED, "triton_*_silu"
        return Op.TRITON_FUSED, "triton_*"

    # SoftMax operations
    if "softmax" in nl:
        return Op.SOFTMAX, "softmax_*"
    
    # Reduce operations - catch more patterns
    if "reduce" in nl:
        if "reduce_segments" in nl or "devicereduce" in nl:
            return Op.REDUCE, "reduce_segments"
        if "reduce_kernel" in nl:
            return Op.REDUCE, "reduce_kernel"
        return Op.REDUCE, "reduce_*"
    
    # Scan operations (library internals - similar to reduce)
    if "scan" in nl and ("cub::" in name or "rocprim::" in name or "at_cuda_detail::cub::" in name):
        return Op.REDUCE, "cub/rocprim::scan_*"

    # Sorting operations (common in sampling/topk)
    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*"

    # Indexing/Scatter/Gather operations
    if any(x in nl for x in ["indices", "scatter", "gather", "index_select", "embedding"]):
        return Op.INDEXING, "index/scatter_*"
    
    # Memory copy operations
    if "copy" in nl or "memcpy" in nl or "_copy_page_indices" in nl:
        return Op.COPY_MEMORY, "copy_*"

    # PyTorch native operations (catch-all for at::native)
    if "at::native::" in name:
        # Try to be more specific
        if "fill" in nl:
            return Op.ELEMENTWISE, "at::native::fill_*"
        return Op.ELEMENTWISE, "at::native::*"

    # ROCm/CUDA library kernels (other)
    if "rocprim::" in name or "cub::" in name:
        return Op.OTHER, "rocprim/cub_*"

    # Fallback: Heuristic fusion detection for unknown Triton kernels
    # If a kernel has multiple operation keywords, it's likely fused
    heuristic_result = _detect_heuristic_fusion(name)
    if heuristic_result:
        return heuristic_result

    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()

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

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

    # KV Cache
    if "reshape_and_cache" in nl:
        return "KV_Cache"

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

    # SwiGLU / Activations
    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"

    # Triton fused operations (generic)
    if "triton_poi" in nl:
        return "Triton_Pointwise"
    if "triton_red" in nl:
        return "Triton_Reduce"
    if "triton_per" in nl:
        return "Triton_Persistent"

    # Reduce operations
    if "reduce" in nl:
        return "Reduce"

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

    # Softmax
    if "softmax" in nl:
        return "Softmax"
    
    # Indexing/Scatter/Gather
    if any(x in nl for x in ["indices", "scatter", "gather", "index_select", "embedding"]):
        return "Indexing"

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

    # Copy/Memory operations
    if "copy" in nl or "memcpy" in nl:
        return "MemCopy"

    return "Other"
