"""
muh_cc_dispatch.py — Unified kernel policy dispatch for BI-V100
================================================================

Python port of CCCL's cc_dispatch.cuh architecture.

CCCL dispatch pattern (cc_dispatch.cuh):
  dispatch_compute_cap(policy_selector, device_cc, functor)
    → policy_getter<PolicySelector, CC>{}()
    → concrete policy struct (ReducePolicy, ScanPolicy, etc.)

Our equivalent:
  dispatch_kernel_config(hardware, kernel_name, **kwargs)
    → policy_for_kernel(kernel_name, hardware, dtype, ...)
    → concrete config dict (threads, items, block_sizes, etc.)

Key insight from cc_dispatch.cuh line 62 (lowest_cc_resolver):
  CCCL collapses architectures with identical policies — if SM80 and SM86
  produce the same ReducePolicy, only one kernel instantiation is generated.
  Our equivalent: pre-compute all configs at import time (see bottom of file)
  so dispatch is a dict lookup, not a function call.

Key insight from dispatch_reduce.cuh line 490:
  dispatch_compute_cap is called ONCE per DeviceReduce invocation.
  The policy is then threaded through InvokeSingleTile / InvokePasses.
  Our equivalent: dispatch_kernel_config returns a frozen config dict
  that's threaded through the entire kernel call chain.

CCCL source files that informed this design:
  cub/detail/cc_dispatch.cuh           — dispatch mechanism
  cub/device/dispatch/dispatch_reduce.cuh   — reduce two-path dispatch
  cub/device/dispatch/dispatch_transform.cuh — transform spread_out_items
  cub/device/dispatch/dispatch_topk.cuh      — topk radix select
  cub/device/dispatch/dispatch_common.cuh    — shared enums
  cub/grid/grid_even_share.cuh               — work distribution
  cub/agent/agent_reduce.cuh                 — tile consumption patterns
  thrust/examples/summary_statistics.cu      — compound reduce pattern
"""

from dataclasses import dataclass, field
from typing import Dict, Optional, Any
import math


# ══════════════════════════════════════════════════════════════
# Hardware descriptor (mirrors muh/include/muh/hardware.cuh)
# ══════════════════════════════════════════════════════════════

@dataclass(frozen=True)
class HardwareCapability:
    """Mirrors muh::hardware_capability from hardware.cuh."""
    vendor: str = "iluvatar"
    arch_version: int = 100
    warp_size: int = 32
    max_threads_per_block: int = 1024
    max_shared_memory_per_block: int = 49152  # 48KB confirmed via ixsmi
    max_registers_per_thread: int = 255
    l2_cache_size_bytes: int = 6 * 1024 * 1024  # 6MB
    memory_bandwidth_gbps: int = 900
    sm_count: int = 16  # CONFIRMED: 16 SMs, NOT 50

    @property
    def bandwidth_per_sm_gbps(self) -> float:
        return self.memory_bandwidth_gbps / self.sm_count

    @property
    def bytes_in_flight(self) -> int:
        """Optimal bytes in flight per SM.
        
        From CCCL tuning_transform.cuh cc_to_min_bytes_in_flight:
          V100=12KB, A100=16KB, H100=48KB, B200=64KB
        BI-V100 per-SM BW = 900/16 = 56 GB/s ≈ B200 level → 64KB
        Confirmed by bench_bi100.py: bif=8 (64KB) wins at all sizes.
        """
        return 64 * 1024

    def at_least(self, vendor: str, min_arch: int) -> bool:
        """Mirrors hardware_capability::at_least()."""
        return self.vendor == vendor and self.arch_version >= min_arch


BI_V100 = HardwareCapability()


# ══════════════════════════════════════════════════════════════
# Kernel configuration structs
# (mirrors CCCL's ReducePolicy, ScanPolicy, TopkPolicy, etc.)
# ══════════════════════════════════════════════════════════════

@dataclass(frozen=True)
class AttentionConfig:
    """Config for paged attention V1/V2 + prefix prefill.
    
    Dispatch axes (from muh_kernel_map.py VLLM_KERNEL_MAP):
      paged_attention_v1: reduce (score reduction per head)
      paged_attention_v2: reduce + scan (partitioned reduce + merge)
      context_attention_fwd: scan + reduce + transform (Triton prefill)
    """
    # Triton flash attention (prefill)
    triton_block_m: int = 32
    triton_block_n: int = 32
    triton_num_warps: int = 4
    triton_num_stages: int = 1
    # Paged attention (decode)
    partition_size: int = 512
    v1_v2_threshold: int = 8192
    # PyTorch fallback (long decode)
    pytorch_decode_threshold: int = 32768
    pytorch_max_tile_blocks: int = 1024
    # Backend selection
    use_native_v1: bool = True
    use_native_v2: bool = False


@dataclass(frozen=True)
class MoEConfig:
    """Config for fused MoE kernel.
    
    Maps to fused_moe_kernel's tl.constexpr parameters.
    Critical for Qwen3.6: 256 experts, top-8, 64 layers.
    """
    block_size_m: int = 64
    block_size_n: int = 64
    block_size_k: int = 32
    group_size_m: int = 8


@dataclass(frozen=True)
class TransformConfig:
    """Config for element-wise ops (SiLU, RMSNorm, RoPE).
    
    From CCCL dispatch_transform.cuh spread_out_items_per_thread:
      items = ceil_div(num_items, sm_count × threads × max_occupancy)
      clamped to [min_items, max_items]
    """
    bytes_in_flight: int = 64 * 1024  # 64KB for BI-V100
    # These are used by ixformer native kernels (not directly tunable)
    # but inform our SMEM budget calculations
    max_smem_per_block: int = 49152


@dataclass(frozen=True)
class CacheConfig:
    """Config for KV cache operations (copy, swap, reshape_and_cache)."""
    copy_block_size: int = 256


# ══════════════════════════════════════════════════════════════
# CCCL-style SMEM constraint checker
# (from muh_kernel_map.py check_smem, used across all policies)
# ══════════════════════════════════════════════════════════════

def check_smem(threads: int, items: int, elem_bytes: int,
               smem_limit: int = BI_V100.max_shared_memory_per_block) -> dict:
    """Verify tile fits in shared memory. Used by all policy selectors."""
    tile_bytes = threads * items * elem_bytes
    max_items = smem_limit // (threads * elem_bytes) if threads * elem_bytes > 0 else 0
    return {
        "tile_bytes": tile_bytes,
        "fits": tile_bytes <= smem_limit,
        "utilization": tile_bytes / smem_limit if smem_limit > 0 else 0,
        "max_items": max_items,
    }


# ══════════════════════════════════════════════════════════════
# CCCL GridEvenShare work distribution (grid_even_share.cuh)
# ══════════════════════════════════════════════════════════════

def grid_even_share(num_items: int, tile_size: int,
                    sm_count: int = BI_V100.sm_count,
                    subscription_factor: int = 5) -> dict:
    """Python port of GridEvenShare::DispatchInit.
    
    CCCL formula: max_blocks = sm_occupancy × sm_count × subscription_factor
    Then items are evenly distributed across blocks, with 'big' blocks
    getting one extra tile.
    """
    if num_items <= 0 or tile_size <= 0:
        return {"grid_size": 0, "total_tiles": 0}

    total_tiles = math.ceil(num_items / tile_size)
    max_grid_size = sm_count * subscription_factor  # ~80 for BI-V100
    grid_size = min(total_tiles, max_grid_size)
    avg_tiles = total_tiles // grid_size if grid_size > 0 else 0
    big_shares = total_tiles - (avg_tiles * grid_size) if grid_size > 0 else 0

    return {
        "grid_size": grid_size,
        "total_tiles": total_tiles,
        "avg_tiles_per_block": avg_tiles,
        "big_shares": big_shares,
        "tile_size": tile_size,
    }


# ══════════════════════════════════════════════════════════════
# Policy selectors (mirrors each algorithm's policy_selector)
# ══════════════════════════════════════════════════════════════

def select_attention_config(
    hw: HardwareCapability,
    dtype_size: int,  # element size in bytes (2=fp16, 4=fp32)
    head_dim: int,
    max_seq_len: int,
    num_kv_heads: int,
) -> AttentionConfig:
    """CCCL-style policy selector for paged attention.
    
    Mirrors: dispatch_reduce.cuh two-path dispatch
      single-tile: num_items ≤ threads × items → V1 (one CTA)
      multi-tile: GridEvenShare → V2 (partitioned + merge)
    
    SMEM constraint for Triton prefill:
      SMEM = BLOCK_N × head_dim × elem_size × 2 (K + V staging)
    """
    smem = hw.max_shared_memory_per_block

    # Triton BLOCK_N: largest that fits SMEM
    triton_block_n = 64
    while triton_block_n * head_dim * dtype_size * 2 > smem and triton_block_n > 16:
        triton_block_n //= 2

    triton_block_m = triton_block_n
    triton_num_warps = 4
    triton_num_stages = 1  # no async copy on BI-V100

    # V1/V2 threshold: V2 worthwhile when partitions > 1 AND
    # per-partition work exceeds merge overhead
    v1_threshold = 8192

    # PyTorch decode threshold: fall back for seq_len > this
    # (ixformer V1 hangs on very long sequences)
    pytorch_threshold = 32768

    # Tile blocks for PyTorch decode: from GridEvenShare
    # max_blocks = sm_count × subscription_factor = 80
    # Each tile processes ~16K tokens (1024 blocks × block_size=16)
    pytorch_tile_blocks = 1024

    return AttentionConfig(
        triton_block_m=triton_block_m,
        triton_block_n=triton_block_n,
        triton_num_warps=triton_num_warps,
        triton_num_stages=triton_num_stages,
        partition_size=512,
        v1_v2_threshold=v1_threshold,
        pytorch_decode_threshold=pytorch_threshold,
        pytorch_max_tile_blocks=pytorch_tile_blocks,
        use_native_v1=True,
        use_native_v2=False,
    )


def select_moe_config(
    hw: HardwareCapability,
    num_experts: int,
    top_k: int,
    hidden_size: int,
    intermediate_size: int,
) -> MoEConfig:
    """Policy selector for fused MoE.
    
    Qwen3.6: 256 experts, top-8, hidden=3584, intermediate=18944
    
    CCCL parallel: each expert is an independent reduce domain.
    With 256 experts × top-8 × batch=1 → 8 active experts per token.
    BI-V100 16 SMs can run 8 expert-matmuls in parallel → one wave.
    """
    # BLOCK_SIZE_M: tokens per tile. For decode (M=1), smallest possible.
    # For prefill (M=4096), larger is better to amortize overhead.
    block_m = 64 if top_k * 1 >= 64 else 32  # decode: top_k tokens
    block_n = 64
    block_k = 32

    # SMEM check: A_tile + B_tile
    # A: block_m × block_k × 2 bytes = 64×32×2 = 4KB
    # B: block_k × block_n × 2 bytes = 32×64×2 = 4KB
    # Total: 8KB << 48KB ✓

    return MoEConfig(
        block_size_m=block_m,
        block_size_n=block_n,
        block_size_k=block_k,
        group_size_m=8,
    )


# ══════════════════════════════════════════════════════════════
# Pre-computed configs (mirrors CCCL compile-time instantiation)
#
# cc_dispatch.cuh line 62: lowest_cc_resolver collapses identical
# policies across CCs. Our equivalent: compute once at import time.
# ══════════════════════════════════════════════════════════════

# Qwen3.6-35B-A3B model parameters (confirmed from qwen3_5.py)
QWEN36_HEAD_DIM = 256       # text_cfg.head_dim
QWEN36_NUM_KV_HEADS = 4     # num_key_value_heads
QWEN36_MAX_SEQ_LEN = 100000 # from computility-run.yaml
QWEN36_MAX_NUM_SEQS = 1     # CRITICAL: computility-run.yaml --max-num-seqs 1
QWEN36_NUM_EXPERTS = 256    # MoE experts
QWEN36_TOP_K = 8            # MoE top-k
QWEN36_HIDDEN = 3584        # hidden_size
QWEN36_INTERMEDIATE = 18944 # intermediate_size

# Pre-computed for fp16 (the common dtype on BI-V100)
ATTENTION_FP16 = select_attention_config(
    hw=BI_V100,
    dtype_size=2,
    head_dim=QWEN36_HEAD_DIM,
    max_seq_len=QWEN36_MAX_SEQ_LEN,
    num_kv_heads=QWEN36_NUM_KV_HEADS,
)

# Pre-computed for bf16
ATTENTION_BF16 = select_attention_config(
    hw=BI_V100,
    dtype_size=2,  # bf16 same size as fp16
    head_dim=QWEN36_HEAD_DIM,
    max_seq_len=QWEN36_MAX_SEQ_LEN,
    num_kv_heads=QWEN36_NUM_KV_HEADS,
)

MOE_CONFIG = select_moe_config(
    hw=BI_V100,
    num_experts=QWEN36_NUM_EXPERTS,
    top_k=QWEN36_TOP_K,
    hidden_size=QWEN36_HIDDEN,
    intermediate_size=QWEN36_INTERMEDIATE,
)

TRANSFORM_CONFIG = TransformConfig(
    bytes_in_flight=BI_V100.bytes_in_flight,
    max_smem_per_block=BI_V100.max_shared_memory_per_block,
)

CACHE_CONFIG = CacheConfig(copy_block_size=256)


# ══════════════════════════════════════════════════════════════
# Unified dispatch entry point
# (mirrors CCCL dispatch_compute_cap)
# ══════════════════════════════════════════════════════════════

_CONFIGS = {
    "attention": ATTENTION_FP16,
    "attention_fp16": ATTENTION_FP16,
    "attention_bf16": ATTENTION_BF16,
    "moe": MOE_CONFIG,
    "transform": TRANSFORM_CONFIG,
    "cache": CACHE_CONFIG,
}


def dispatch_kernel_config(kernel_name: str,
                           hw: HardwareCapability = BI_V100) -> Any:
    """Unified policy dispatch — Python equivalent of dispatch_compute_cap.
    
    Usage:
        config = dispatch_kernel_config("attention")
        # config.triton_block_m, config.partition_size, etc.
        
        config = dispatch_kernel_config("moe")
        # config.block_size_m, config.block_size_n, etc.
    """
    if kernel_name not in _CONFIGS:
        raise KeyError(
            f"Unknown kernel: {kernel_name}. "
            f"Available: {list(_CONFIGS.keys())}"
        )
    return _CONFIGS[kernel_name]


# ══════════════════════════════════════════════════════════════
# CLI: dump all configs for inspection
# ══════════════════════════════════════════════════════════════

if __name__ == "__main__":
    print("muh_cc_dispatch: CCCL-style unified kernel policy dispatch\n")
    print(f"Hardware: {BI_V100.vendor} BI-V100")
    print(f"  SMs: {BI_V100.sm_count}, SMEM: {BI_V100.max_shared_memory_per_block//1024}KB, "
          f"BW: {BI_V100.memory_bandwidth_gbps}GB/s, "
          f"BW/SM: {BI_V100.bandwidth_per_sm_gbps:.1f}GB/s")
    print(f"  bytes_in_flight: {BI_V100.bytes_in_flight//1024}KB\n")

    for name, config in _CONFIGS.items():
        print(f"[{name}]")
        for k, v in config.__dict__.items():
            if not k.startswith('_'):
                print(f"  {k}: {v}")
        print()

    # GridEvenShare example for decode
    print("GridEvenShare example (50K token decode, block_size=16):")
    es = grid_even_share(50000 // 16, tile_size=1024)
    for k, v in es.items():
        print(f"  {k}: {v}")