#!/usr/bin/env python3
"""What can we learn about expert importance WITHOUT a forward pass?

A proper REAP prune needs router-weighted activation norms from calibration data, which
needs an engine that can run deepseek_v41 -- and none exists yet. So before committing to
building one, check whether the checkpoint already carries a usable signal:

  ffn.gate.bias      the noaux_tc load-balancing correction. It is ADDED to the router
                     score purely for top-k selection. Training drives it DOWN for experts
                     the router over-picks and UP for ones it under-picks, so it is a proxy
                     for natural selection frequency -- inverted.
  ffn.gate.bias_vl   a SECOND bias, used for vision-language tokens. If it differs from
                     the text bias, expert specialisation is modality-dependent and a
                     text-only calibration would prune the vision path.
  ffn.gate.weight    per-expert router directions [n_experts, hidden]. Cheap to cluster:
                     experts with near-identical directions are selected by similar tokens
                     and are merge candidates. This is the data-free fallback.
"""
import json, os, sys, glob
import torch
from safetensors import safe_open

src = sys.argv[1]
idx = json.load(open(os.path.join(src, "model.safetensors.index.json")))["weight_map"]

layers = sorted({int(k.split(".")[1]) for k in idx if k.endswith("ffn.gate.bias")})
print(f"MoE layers with a router bias: {len(layers)} -> {layers[:5]}...{layers[-3:]}\n")

open_files = {}
def get(name):
    sh = idx[name]
    if sh not in open_files:
        open_files[sh] = safe_open(os.path.join(src, sh), framework="pt")
    return open_files[sh].get_tensor(name)

print(f"{'layer':>5} {'n_exp':>6} {'bias std':>9} {'bias range':>16} "
      f"{'|txt-vl| mean':>13} {'corr(txt,vl)':>12} {'top/bot 10% gap':>15}")
rows = []
for L in layers:
    b = get(f"layers.{L}.ffn.gate.bias").float()
    bvl = get(f"layers.{L}.ffn.gate.bias_vl").float()
    n = b.numel()
    srt = b.sort().values
    k = max(1, n // 10)
    gap = (srt[-k:].mean() - srt[:k].mean()).item()
    corr = torch.corrcoef(torch.stack([b, bvl]))[0, 1].item()
    rows.append((L, n, b.std().item(), b.min().item(), b.max().item(),
                 (b - bvl).abs().mean().item(), corr, gap))
    if L in layers[:3] + layers[len(layers)//2:len(layers)//2+1] + layers[-2:]:
        print(f"{L:>5} {n:>6} {b.std():>9.4f} [{b.min():>6.3f},{b.max():>6.3f}] "
              f"{(b-bvl).abs().mean():>13.4f} {corr:>12.4f} {gap:>15.4f}")

import statistics
print(f"\nacross all {len(rows)} layers:")
print(f"  bias std           mean {statistics.mean(r[2] for r in rows):.4f}")
print(f"  corr(text, vl)     mean {statistics.mean(r[6] for r in rows):.4f} "
      f"min {min(r[6] for r in rows):.4f}")
print(f"  |text - vl| bias   mean {statistics.mean(r[5] for r in rows):.4f}")

# router-direction redundancy: how many experts are near-duplicates?
print("\nRouter-direction cosine similarity (the data-free merge signal):")
for L in layers[:2] + layers[len(layers)//2:len(layers)//2+1] + layers[-1:]:
    W = get(f"layers.{L}.ffn.gate.weight").float()
    Wn = W / W.norm(dim=1, keepdim=True).clamp(min=1e-9)
    C = Wn @ Wn.T
    C.fill_diagonal_(-1.0)
    best = C.max(dim=1).values
    print(f"  layer {L:>2}: nearest-neighbour cosine  mean {best.mean():.4f}  "
          f"max {best.max():.4f}  >0.9: {int((best>0.9).sum())}/{W.shape[0]}  "
          f">0.7: {int((best>0.7).sum())}")