#!/usr/bin/env python3
"""Diagnostic: track per-layer magnitudes to find residual explosion.

Runs single token "The" through all 61 layers and prints:
  |X_in|, |x_normed|, |F_attn|, |X_mid|, |F_ffn|, |X_next|
  for each layer.

This identifies WHERE the residual stream starts exploding.

Usage (on B200):
    source /root/dsv4-nvfp4-workspace/venv/bin/activate
    cd /root/dsv4-nvfp4-workspace/kernel
    python3 tests/test_residual_diagnostic.py
"""
import os, sys, math, json
import torch
from pathlib import Path

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
NUM_GPUS = 8

# Import helpers from the main test
from tests.test_minimal_e2e import (
    FP4_LUT, dequant_nvfp4_weight, nvfp4_linear, RMSNorm,
    build_rope_cache, apply_rope_partial, apply_inverse_rope,
    load_weights_to_cpu, get_layer_weights
)
from single_shot_inference import mHCBlock


def main():
    device = 'cuda:0'
    with open(os.path.join(CHECKPOINT_DIR, "config.json")) as f:
        cfg = json.load(f)

    n_layers = cfg["num_hidden_layers"]
    H = cfg["hidden_size"]
    n_h = cfg["num_attention_heads"]
    hd = cfg["head_dim"]
    rd = cfg.get("qk_rope_head_dim", cfg.get("rope_dim", 64))
    n_hc = 4
    o_rank = cfg.get("output_group_dim", 1024)
    o_groups = cfg.get("num_output_groups", 16)
    heads_per_group = n_h // o_groups
    group_input_dim = heads_per_group * hd

    print("Loading weights to CPU...")
    all_weights = load_weights_to_cpu(CHECKPOINT_DIR)

    embed_w = all_weights.get("model.embed_tokens.weight")
    embed = torch.nn.Embedding.from_pretrained(embed_w.bfloat16().to(device))

    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR)
    input_ids = tokenizer.encode("The")
    tid = torch.tensor([input_ids[0]], dtype=torch.long, device=device)
    positions = torch.tensor([0], dtype=torch.long, device=device)

    emb = embed(tid)
    from dsv4.layers.mhc import mHCLayer
    X = mHCLayer.init_state(emb, n_hc)

    print(f"\n{'L':>3} {'|X_in|':>10} {'|x_norm|':>10} {'|F_attn|':>10} {'|X_mid|':>10} {'|F_ffn|':>10} {'|X_out|':>10}  nan? inf?")

    for li in range(n_layers):
        gpu = li % NUM_GPUS
        dev = f"cuda:{gpu}"
        X = X.to(dev)
        torch.cuda.set_device(gpu)

        w = get_layer_weights(all_weights, li, dev)
        rope_cos, rope_sin = build_rope_cache(8192, rd, dev)
        positions_dev = positions.to(dev)

        # Build per-layer components
        attn_mhc = mHCBlock(hidden_dim=H, n_hc=n_hc, device=dev)
        attn_mhc.load_from_checkpoint(
            all_weights[f"model.layers.{li}.attn_hc.fn"],
            all_weights[f"model.layers.{li}.attn_hc.base"],
            all_weights[f"model.layers.{li}.attn_hc.scale"])
        ffn_mhc = mHCBlock(hidden_dim=H, n_hc=n_hc, device=dev)
        ffn_mhc.load_from_checkpoint(
            all_weights[f"model.layers.{li}.ffn_hc.fn"],
            all_weights[f"model.layers.{li}.ffn_hc.base"],
            all_weights[f"model.layers.{li}.ffn_hc.scale"])
        attn_norm = RMSNorm(H, eps=cfg.get('rms_norm_eps', 1e-6), device=dev)
        attn_norm.weight = all_weights[f"model.layers.{li}.input_layernorm.weight"].to(device=dev, dtype=torch.float32)
        ffn_norm = RMSNorm(H, eps=cfg.get('rms_norm_eps', 1e-6), device=dev)
        ffn_norm.weight = all_weights[f"model.layers.{li}.post_attention_layernorm.weight"].to(device=dev, dtype=torch.float32)

        # Track magnitudes
        x_in_mag = X.float().abs().max().item()

        # ATTENTION
        x_in_attn, attn_ctx = attn_mhc.pre_block(X)
        x_normed = attn_norm.forward(x_in_attn)
        x_normed_mag = x_normed.float().abs().max().item()

        pre = f"model.layers.{li}.self_attn"

        # Q projection
        c_Q = nvfp4_linear(x_normed, w[f"{pre}.q_a_proj.weight"],
                           w[f"{pre}.q_a_proj.weight_scale"], w[f"{pre}.q_a_proj.weight_scale_2"])
        q_norm_w = w.get(f"{pre}.q_a_norm.weight")
        if q_norm_w is not None:
            c_Q_f = c_Q.float()
            c_Q = (c_Q_f * c_Q_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() * q_norm_w.float()).bfloat16()
        q = nvfp4_linear(c_Q, w[f"{pre}.q_b_proj.weight"],
                         w[f"{pre}.q_b_proj.weight_scale"], w[f"{pre}.q_b_proj.weight_scale_2"])

        kv = nvfp4_linear(x_normed, w[f"{pre}.kv_proj.weight"],
                          w[f"{pre}.kv_proj.weight_scale"], w[f"{pre}.kv_proj.weight_scale_2"])
        kv_norm_w = w.get(f"{pre}.kv_norm.weight")
        if kv_norm_w is not None:
            kv_f = kv.float()
            kv = (kv_f * kv_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() * kv_norm_w.float()).bfloat16()

        q_heads = q.reshape(1, n_h, hd)
        kv_new = kv.reshape(1, 1, hd)
        q_roped = apply_rope_partial(q_heads, positions_dev, rope_cos, rope_sin, hd, rd)
        kv_roped = apply_rope_partial(kv_new, positions_dev, rope_cos, rope_sin, hd, rd)

        k_exp = kv_roped.expand(n_h, -1, -1).contiguous()
        v_exp = kv_roped.expand(n_h, -1, -1).contiguous()
        q_input = q_roped.permute(1, 0, 2)
        scale = 1.0 / math.sqrt(hd)

        sink_key = f"{pre}.sinks"
        if sink_key in w:
            sinks = w[sink_key].to(device=dev)
            sink_k = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device=dev)
            sink_v = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device=dev)
            k_ws = torch.cat([k_exp, sink_k], dim=1)
            v_ws = torch.cat([v_exp, sink_v], dim=1)
            sink_bias = torch.zeros(n_h, 1, 2, dtype=torch.bfloat16, device=dev)
            for h in range(n_h):
                sink_bias[h, :, -1] = sinks[h]
            attn_out = torch.nn.functional.scaled_dot_product_attention(
                q_input, k_ws, v_ws, attn_mask=sink_bias, scale=scale)
        else:
            attn_out = torch.nn.functional.scaled_dot_product_attention(
                q_input, k_exp, v_exp, scale=scale, is_causal=False)
        attn_out = attn_out.permute(1, 0, 2)

        attn_out = apply_inverse_rope(attn_out, positions_dev, rope_cos, rope_sin, hd, rd)

        # Output projection
        attn_flat = attn_out.reshape(1, n_h * hd)
        attn_grouped = attn_flat.reshape(1, o_groups, heads_per_group * hd)
        oa_w = w[f"{pre}.o_a_proj.weight"].bfloat16()
        oa_3d = oa_w.reshape(o_groups, o_rank, group_input_dim)
        grouped_out = torch.bmm(attn_grouped.permute(1, 0, 2), oa_3d.transpose(1, 2))
        grouped_flat = grouped_out.permute(1, 0, 2).reshape(1, o_groups * o_rank)
        F_attn = nvfp4_linear(grouped_flat, w[f"{pre}.o_b_proj.weight"],
                               w[f"{pre}.o_b_proj.weight_scale"], w[f"{pre}.o_b_proj.weight_scale_2"])
        f_attn_mag = F_attn.float().abs().max().item()

        X_mid = attn_mhc.post_block(X, F_attn, attn_ctx)
        x_mid_mag = X_mid.float().abs().max().item()

        # FFN (shared expert + routed, using the FIXED MoE loop)
        x_ffn, ffn_ctx = ffn_mhc.pre_block(X_mid)
        x_ffn_normed = ffn_norm.forward(x_ffn)

        # Routed MoE
        n_experts = cfg["n_routed_experts"]
        top_k = cfg.get("num_experts_per_tok", 6)
        routed_scaling = cfg.get("routed_scaling_factor", 2.5)
        swiglu_limit = cfg.get("swiglu_limit", 10.0)
        is_hash = li < 3

        if is_hash:
            tid2eid_key = f"model.layers.{li}.mlp.gate.tid2eid"
            if tid2eid_key in w:
                tid2eid = w[tid2eid_key]
                tid_val = tid.item() if tid.device == dev else tid.to(dev).item()
                expert_ids = tid2eid[tid_val]
                expert_weights = torch.ones(top_k, dtype=torch.float32, device=dev) / top_k
            else:
                is_hash = False

        if not is_hash:
            gate_w = w[f"model.layers.{li}.mlp.gate.weight"]
            logits = torch.nn.functional.linear(x_ffn_normed, gate_w.bfloat16())
            activated = torch.sqrt(torch.nn.functional.softplus(logits.float()) + 1e-6)
            e_bias_key = f"model.layers.{li}.mlp.gate.e_bias"
            if e_bias_key in w:
                activated = activated + w[e_bias_key].float().unsqueeze(0)
            scores, indices = activated.topk(top_k, dim=-1)
            unbiased = torch.sqrt(torch.nn.functional.softplus(logits.float()) + 1e-6)
            unbiased_scores = torch.gather(unbiased, -1, indices)
            expert_weights = unbiased_scores / unbiased_scores.sum(dim=-1, keepdim=True)
            expert_ids = indices[0]

        expert_outputs = []
        for i, eid in enumerate(expert_ids):
            eid_int = eid.item()
            epre = f"model.layers.{li}.mlp.experts.{eid_int}"
            gate = nvfp4_linear(x_ffn_normed, w[f"{epre}.gate_proj.weight"],
                               w[f"{epre}.gate_proj.weight_scale"], w[f"{epre}.gate_proj.weight_scale_2"])
            up = nvfp4_linear(x_ffn_normed, w[f"{epre}.up_proj.weight"],
                             w[f"{epre}.up_proj.weight_scale"], w[f"{epre}.up_proj.weight_scale_2"])
            silu_out = torch.nn.functional.silu(gate.float())
            if swiglu_limit is not None:
                silu_out = silu_out.clamp(-swiglu_limit, swiglu_limit)
                up_clamped = up.float().clamp(-swiglu_limit, swiglu_limit)
            else:
                up_clamped = up.float()
            hidden = (silu_out * up_clamped).bfloat16()
            down = nvfp4_linear(hidden, w[f"{epre}.down_proj.weight"],
                               w[f"{epre}.down_proj.weight_scale"], w[f"{epre}.down_proj.weight_scale_2"])
            expert_outputs.append(down)

        routed_out = torch.zeros_like(x_ffn_normed)
        for i, (out, wt) in enumerate(zip(expert_outputs, expert_weights)):
            w_val = wt.item() if wt.dim() == 0 else wt[i].item() if wt.dim() == 1 else wt.flatten()[i].item()
            routed_out = routed_out + (out.float() * w_val).bfloat16()
        routed_out = (routed_out.float() * routed_scaling).bfloat16()

        # Shared expert
        se_pre = f"model.layers.{li}.mlp.shared_experts"
        gate = nvfp4_linear(x_ffn_normed, w[f"{se_pre}.gate_proj.weight"],
                           w[f"{se_pre}.gate_proj.weight_scale"], w[f"{se_pre}.gate_proj.weight_scale_2"])
        up = nvfp4_linear(x_ffn_normed, w[f"{se_pre}.up_proj.weight"],
                         w[f"{se_pre}.up_proj.weight_scale"], w[f"{se_pre}.up_proj.weight_scale_2"])
        hidden = (torch.nn.functional.silu(gate.float()).clamp(-10, 10) * up.float().clamp(-10, 10)).bfloat16()
        shared_out = nvfp4_linear(hidden, w[f"{se_pre}.down_proj.weight"],
                                  w[f"{se_pre}.down_proj.weight_scale"], w[f"{se_pre}.down_proj.weight_scale_2"])

        F_ffn = routed_out + shared_out
        f_ffn_mag = F_ffn.float().abs().max().item()

        X_next = ffn_mhc.post_block(X_mid, F_ffn, ffn_ctx)
        x_out_mag = X_next.float().abs().max().item()

        has_nan = torch.isnan(X_next).any().item()
        has_inf = torch.isinf(X_next).any().item()

        print(f"{li:3d} {x_in_mag:10.3f} {x_normed_mag:10.3f} {f_attn_mag:10.3f} {x_mid_mag:10.3f} {f_ffn_mag:10.3f} {x_out_mag:10.3f}  {'NaN' if has_nan else ''} {'INF' if has_inf else ''}")

        X = X_next
        del w
        torch.cuda.empty_cache()

    # Final logits
    X = X.to('cuda:0')
    lm_w = all_weights.get("lm_head.weight", embed_w).bfloat16().to('cuda:0')
    final_norm_w = all_weights.get("model.norm.weight")
    x_out = X[:, 0, :]
    if final_norm_w is not None:
        xf = x_out.float()
        rms = xf.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt()
        x_out = (xf * rms * final_norm_w.to('cuda:0').float()).bfloat16()
    logits = torch.nn.functional.linear(x_out, lm_w)
    top5_vals, top5_ids = torch.topk(logits[0], 5)
    top5_str = ' '.join([f'{tokenizer.decode([t.item()])}({v.item():.1f})'
                         for t, v in zip(top5_ids, top5_vals)])
    print(f"\nTop-5: {top5_str}")


if __name__ == "__main__":
    main()
