"""D-57 importance weights for the N-bit scale search: llama.cpp's imatrix (mean squared input activation per input
channel, per expert) for each layer's experts, as the trainer and exporter read them.

`scripts/imatrix_to_importance.py` turns llama.cpp's imatrix.gguf into a safetensors file with, per layer i,
`layers.{i}.gate_up` [E, hidden] (gate and up read the same input) and `layers.{i}.down` [E, moe_intermediate] in fp32.
An expert the calibration text never routed to (count 0) gets its layer's mean over the routed experts.
"""
from __future__ import annotations

from pathlib import Path

import numpy as np
import torch


def from_imatrix(gguf_path: str | Path, gguf) -> dict[str, torch.Tensor]:
    """llama.cpp imatrix.gguf → {"layers.{i}.gate_up": [E, in], "layers.{i}.down": [E, in]} (fp32)."""
    r = gguf.GGUFReader(str(gguf_path))
    by = {t.name: np.asarray(t.data, dtype=np.float32) for t in r.tensors}
    out = {}
    i = 0
    while f"blk.{i}.ffn_gate_exps.weight.in_sum2" in by:
        for ours, theirs in (("gate_up", "ffn_gate_exps"), ("down", "ffn_down_exps")):
            s2 = by[f"blk.{i}.{theirs}.weight.in_sum2"]            # [E, in]
            n = by[f"blk.{i}.{theirs}.weight.counts"].reshape(-1, 1)  # [E, 1]
            seen = n[:, 0] > 0
            imp = np.where(seen[:, None], s2 / np.maximum(n, 1), 0.0)
            imp[~seen] = imp[seen].mean(0) if seen.any() else 1.0
            out[f"layers.{i}.{ours}"] = torch.from_numpy(np.ascontiguousarray(imp, dtype=np.float32))
        i += 1
    return out


def load(path: str | Path) -> dict[str, torch.Tensor]:
    from safetensors.torch import load_file
    return load_file(str(path))