#!/usr/bin/env python3
"""
Layer 0 kernel comparison test: CuTeDSL NVFP4 kernel vs BF16 reference.

No vLLM, no Docker, no tensor parallelism. Just raw weights + CuTeDSL kernel.
If cosine < 0.99, the test exits with error.

Uses the bridge layer in cutedsl/bridge.py for tensor layout conversion.
"""

import os
import sys
import json
import glob
import torch
from safetensors import safe_open

# Add repo root to path
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, REPO_ROOT)

from cutedsl.bridge import (
    quantize_to_nvfp4,
    quantize_weight_to_nvfp4,
    assemble_scales_2d_side,
    assemble_scales_3d_side,
    make_b_k_major,
    compute_expert_offsets,
    run_nvfp4_grouped_gemm,
)

from cutedsl.moe_pipeline import (
    stage_activation,
    prepare_nvfp4_moe_weights,
    run_nvfp4_moe,
)

# ── Constants ──────────────────────────────────────────────────────────

NVFP4_MODEL_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
LAYER_IDX = 0
DEVICE = "cuda"
COSINE_THRESHOLD = 0.99

# E2M1 FP4 lookup table (for BF16 dequant reference)
E2M1_LUT = torch.tensor([
    0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
    -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
], dtype=torch.float32)


# ── Checkpoint loading ─────────────────────────────────────────────────

def find_shards(model_dir):
    index_path = os.path.join(model_dir, "model.safetensors.index.json")
    key_to_shard = {}
    if os.path.exists(index_path):
        with open(index_path) as f:
            index = json.load(f)
        for key, shard in index["weight_map"].items():
            key_to_shard[key] = os.path.join(model_dir, shard)
    else:
        for sf in glob.glob(os.path.join(model_dir, "*.safetensors")):
            with safe_open(sf, framework="pt") as f:
                for key in f.keys():
                    key_to_shard[key] = sf
    return key_to_shard


def load_layer_tensors(model_dir, layer_idx):
    """Load all tensors for a specific layer. Keys normalized (no 'model.' prefix)."""
    key_to_shard = find_shards(model_dir)
    layer_prefix = f"layers.{layer_idx}."
    shard_to_keys = {}
    for key, shard in key_to_shard.items():
        norm_key = key.removeprefix("model.")
        if not norm_key.startswith(layer_prefix):
            continue
        shard_to_keys.setdefault(shard, []).append((key, norm_key))
    tensors = {}
    for shard, keys in shard_to_keys.items():
        with safe_open(shard, framework="pt") as f:
            for orig_key, norm_key in keys:
                tensors[norm_key] = f.get_tensor(orig_key)
    return tensors


# ── NVFP4 Dequantization (BF16 reference) ─────────────────────────────

def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale):
    """Dequantize NVFP4 (E2M1 + E4M3 + global) to BF16."""
    device = packed_uint8.device
    lut = E2M1_LUT.to(device)
    lower = lut[(packed_uint8 & 0x0F).long()]
    upper = lut[((packed_uint8 >> 4) & 0x0F).long()]
    out_features = packed_uint8.shape[0]
    in_features = packed_uint8.shape[1] * 2
    unpacked = torch.empty(out_features, in_features, dtype=torch.float32, device=device)
    unpacked[:, 0::2] = lower
    unpacked[:, 1::2] = upper
    block_scale = scale_e4m3.float()
    block_expanded = block_scale.repeat_interleave(16, dim=1)[:, :in_features]
    return (unpacked * block_expanded * global_scale).to(torch.bfloat16)


def dequantize_nvfp4_experts(nvfp4_tensors, layer_idx, expert_indices):
    """Dequantize expert weights from NVFP4 checkpoint → BF16."""
    experts = {}
    for e in expert_indices:
        expert = {}
        for proj in ["gate_proj", "up_proj", "down_proj"]:
            weight_key = f"layers.{layer_idx}.mlp.experts.{e}.{proj}.weight"
            scale_key = f"layers.{layer_idx}.mlp.experts.{e}.{proj}.weight_scale"
            gs_key = f"layers.{layer_idx}.mlp.experts.{e}.{proj}.weight_scale_2"
            if weight_key not in nvfp4_tensors:
                if proj == "down_proj" and e == 211:
                    continue
                raise KeyError(f"Missing {weight_key}")
            weight = nvfp4_tensors[weight_key].to(DEVICE)
            scale = nvfp4_tensors[scale_key].to(DEVICE)
            global_scale = nvfp4_tensors[gs_key].item()
            expert[proj] = dequantize_nvfp4_weight(weight, scale, global_scale)
        experts[e] = expert
    return experts


# ── BF16 MoE Forward ───────────────────────────────────────────────────

def moe_forward_bf16(hidden_states, experts, expert_ids, expert_weights):
    """Run MoE forward pass in pure BF16 (torch.matmul)."""
    num_tokens, hidden_size = hidden_states.shape
    top_k = expert_ids.shape[1]
    output = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE)
    for t in range(num_tokens):
        for k in range(top_k):
            e = expert_ids[t, k].item()
            w = expert_weights[t, k].item()
            if e not in experts:
                continue
            x = hidden_states[t]
            gate = x @ experts[e]["gate_proj"].T
            up = x @ experts[e]["up_proj"].T
            activated = torch.nn.functional.silu(gate) * up
            if "down_proj" in experts[e]:
                y = activated @ experts[e]["down_proj"].T
            else:
                y = activated[:hidden_size]
            output[t] += w * y
    return output


# ── CuTeDSL NVFP4 Kernel MoE Forward ──────────────────────────────────

def moe_forward_nvfp4_l1_only(slot_hidden, nvfp4_tensors, layer_idx, expert_indices, tokens_per_expert):
    """Run L1 (gate+up) GEMM using CuTeDSL.
    
    slot_hidden is already laid out slot-major: [expert0_tokens | expert1_tokens | ...]
    """
    num_slots, hidden_size = slot_hidden.shape
    num_experts = len(expert_indices)

    # Quantize activation
    x_fp4, x_sf, x_igs = quantize_to_nvfp4(slot_hidden)

    # Load and quantize weights
    w_fp4_list = []
    w_sf_list = []
    w_gs_list = []

    for e in expert_indices:
        gate_w_key = f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight"
        gate_sf_key = f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale"
        gate_gs_key = f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale_2"
        up_w_key = f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight"
        up_sf_key = f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale"
        up_gs_key = f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale_2"

        gate_w_bf16 = dequantize_nvfp4_weight(
            nvfp4_tensors[gate_w_key].to(DEVICE),
            nvfp4_tensors[gate_sf_key].to(DEVICE),
            nvfp4_tensors[gate_gs_key].item(),
        )
        up_w_bf16 = dequantize_nvfp4_weight(
            nvfp4_tensors[up_w_key].to(DEVICE),
            nvfp4_tensors[up_sf_key].to(DEVICE),
            nvfp4_tensors[up_gs_key].item(),
        )

        # Fuse gate + up, transpose to (K=hidden, N=6144)
        fused = torch.cat([gate_w_bf16, up_w_bf16], dim=0)  # (6144, 7168)
        l1_w_bf16 = fused.T  # (7168, 6144)
        l1_w_fp4, l1_w_sf, l1_w_gs = quantize_weight_to_nvfp4(l1_w_bf16)

        w_fp4_list.append(l1_w_fp4)
        w_sf_list.append(l1_w_sf)
        w_gs_list.append(l1_w_gs)

    # Stack and convert to K-major
    mat_b = make_b_k_major(torch.stack(w_fp4_list))

    # Assemble scale factors
    # scale_a: per-expert activation scales, split by expert offsets
    x_sf_parts = []
    offset = 0
    for tpe in tokens_per_expert:
        x_sf_parts.append(x_sf[offset:offset+tpe])
        offset += tpe
    scale_a = assemble_scales_2d_side(x_sf_parts)
    scale_b = assemble_scales_3d_side(w_sf_list)

    # Expert offsets
    expert_offsets = compute_expert_offsets(tokens_per_expert, num_experts)

    # Global scales
    global_scale_a = torch.tensor([x_igs] * num_experts, dtype=torch.float32, device=DEVICE)
    global_scale_b = torch.tensor(w_gs_list, dtype=torch.float32, device=DEVICE)

    # Run kernel
    out = run_nvfp4_grouped_gemm(
        mat_a=x_fp4, mat_b=mat_b,
        scale_a=scale_a, scale_b=scale_b,
        expert_offsets=expert_offsets,
        global_scale_a=global_scale_a, global_scale_b=global_scale_b,
    )
    return out

def main():
    torch.manual_seed(42)
    expert_indices = [0, 1, 2]
    top_k = 2
    num_tokens = 4
    hidden_size = 7168

    # ── Load NVFP4 checkpoint ──
    print("=" * 70)
    print("  Loading NVFP4 checkpoint layer 0")
    print("=" * 70)

    nvfp4_tensors = load_layer_tensors(NVFP4_MODEL_DIR, LAYER_IDX)
    expert_keys = [k for k in sorted(nvfp4_tensors.keys()) if 'experts.0.' in k]
    print(f"  {len(nvfp4_tensors)} tensors loaded")
    for key in expert_keys[:3]:
        t = nvfp4_tensors[key]
        print(f"    {key}: dtype={t.dtype} shape={tuple(t.shape)}")

    # ── Prepare NVFP4 weights ──
    print("
  Preparing NVFP4 weights (dequant → re-quant)...")
    weights = prepare_nvfp4_moe_weights(nvfp4_tensors, LAYER_IDX, expert_indices)
    print(f"    L1: {len(weights['l1_fp4'])} experts, shape {weights['l1_fp4'][0].shape}")
    print(f"    L2: {len(weights['l2_fp4'])} experts, shape {weights['l2_fp4'][0].shape}")

    # ── Dequantize → BF16 reference ──
    print("
  Dequantizing NVFP4 → BF16 reference...")
    nvfp4_experts_bf16 = dequantize_nvfp4_experts(nvfp4_tensors, LAYER_IDX, expert_indices)

    # ── Create test input ──
    hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) * 2.0
    expert_ids = torch.tensor([[0, 1]] * num_tokens, dtype=torch.int32, device=DEVICE)
    expert_weights = torch.tensor([[0.6, 0.4]] * num_tokens, dtype=torch.float32, device=DEVICE)

    # ── BF16 full MoE reference ──
    print("
  Running BF16 MoE reference...")
    ref_output = moe_forward_bf16(hidden_states, nvfp4_experts_bf16, expert_ids, expert_weights)
    print(f"    BF16 ref: amax={ref_output.abs().max():.4f} mean={ref_output.float().mean():.6f}")

    del nvfp4_experts_bf16
    torch.cuda.empty_cache()

    # ── CuTeDSL NVFP4 full MoE pipeline ──
    print("
  Running CuTeDSL NVFP4 MoE pipeline (first run compiles, ~1-2 min)...")
    kernel_output = run_nvfp4_moe(
        hidden_states, expert_ids, expert_weights,
        weights, expert_indices,
    )
    print(f"    Kernel:   amax={kernel_output.abs().max():.4f} mean={kernel_output.float().mean():.6f}")

    # ── Compare ──
    cosine = torch.nn.functional.cosine_similarity(
        kernel_output.flatten().unsqueeze(0).float(),
        ref_output.flatten().unsqueeze(0).float(),
    ).item()
    mse = (kernel_output.float() - ref_output.float()).pow(2).mean().item()

    print(f"
{'=' * 70}")
    print(f"  RESULT: cosine={cosine:.6f}  MSE={mse:.6e}")
    print(f"{'=' * 70}")

    if cosine < COSINE_THRESHOLD:
        print(f"  ❌ FAIL: cosine {cosine:.6f} < {COSINE_THRESHOLD}")
        sys.exit(1)
    else:
        print(f"  ✅ PASS: cosine {cosine:.6f} >= {COSINE_THRESHOLD}")


if __name__ == "__main__":
    main()