#!/usr/bin/env python3
"""Build per-tensor feature matrix for Gnom training (deterministic).

    build_features.py --model M.gguf --imatrix I.gguf --out gnom/features/<m>.npz
                      [--ssim-table models/ssim_table.npz] [--n-layers N]

Output .npz: X (n_tensors x 16, RAW, see gnom/README), names (tensor names),
stats (timp_mean/std, logN_median for per-model standardization at train).
16 = 6 type one-hot + layer + logN + timp_mean + timp_max + frag + econc
+ 4 weight stats (std/kurt/max-mean/outlier-fraction). The README's older
"15-dim" counts predate the outlier-fraction column (RL-PLAN step 2).
Weight stats are computed chunk-wise from the mmap (bounded RAM) and BF16
tensors are byte-viewed through ml_dtypes (scar 2026-10-16: the old
np.array(t.data, float32) path read BF16 bytes as numbers — 4 garbage
columns on every BF16 model, and the 18 GB Ornith build OOMed on a
float64 copy).

--ssim-table is DEAD (2026-10-16): 4/4 structural metrics rank together and
orthogonal to KLD/PPL — kept only to reproduce archival features byte-for-byte.
"""
import argparse
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import ml_dtypes  # noqa: F401  (bfloat16 decode)
from gguf import GGUFReader
from llama_inspect import read_model
from imatrix_reader import read_imatrix, build_importance_table
from tinynet import featurize


def load_ssim(path):
    if not path or not os.path.exists(path):
        return {}
    try:
        z = np.load(path, allow_pickle=False)
        return {n: float(np.ravel(z[n])[0]) for n in z.files}  # Q4_K first
    except Exception:
        return {}


def _w_view(t):
    """Reader tensor -> typed array view, or None if not decodable.
    gguf yields F16/F32/F64 natively; BF16 arrives as raw uint8 bytes
    (rows x bytes) and must be byte-viewed, never cast."""
    arr = t.data
    if arr.dtype in (np.float16, np.float32, np.float64):
        return arr
    if arr.dtype == np.uint8 and arr.ndim >= 1 and arr.shape[-1] % 2 == 0:
        try:
            return arr.view(ml_dtypes.bfloat16)
        except Exception:
            return None
    return None


def w_stats_stream(arr, chunk_rows=1024):
    """The 4 weight-stat columns of featurize, computed in bounded RAM.
    Two chunked passes; accumulation in float64 scalars over small slabs.
    Formulas are exactly featurize's: log10 std, clipped kurt/10,
    log10(max|w|/mean|w|+1), fraction(|w| > 6*std).
    (skewness was trialed and REMOVED 2026-10-18: zero effect on v4-tz
    0.575, directionally negative on v2-raw, and it switched the best
    variant under noise — destabilizing, not neutral.)"""
    two_d = arr.ndim > 1
    rows = arr.shape[0] if two_d else 1
    step = max(1, int(chunk_rows))
    m1 = m2 = m3 = m4 = amax = asum = 0.0
    nn = 0
    for i in range(0, rows, step):
        blk = np.asarray(arr[i:i + step] if two_d else arr, dtype=np.float64)
        blk = blk.ravel()
        blk2 = blk * blk
        m1 += float(blk.sum())
        m2 += float(blk2.sum())
        m3 += float((blk2 * blk).sum())
        m4 += float((blk2 * blk2).sum())
        a = np.abs(blk)
        am = float(a.max())
        if am > amax:
            amax = am
        asum += float(a.sum())
        nn += blk.size
    if nn == 0:
        return None
    mean = m1 / nn
    var = max(m2 / nn - mean * mean, 0.0)
    std = var ** 0.5
    m4c = m4 / nn - 4.0 * mean * (m3 / nn) + 6.0 * mean * mean * (m2 / nn) \
        - 3.0 * mean ** 4
    kurt = m4c / (var * var + 1e-12) - 3.0
    mxm = amax / (asum / nn + 1e-12)
    thr = 6.0 * (std + 1e-12)
    cnt = 0
    for i in range(0, rows, step):
        blk = np.asarray(arr[i:i + step] if two_d else arr, dtype=np.float32)
        cnt += int(np.count_nonzero(np.abs(blk) > thr))
    outl = cnt / nn
    return (float(np.log10(std + 1e-12)),
            float(np.clip(kurt, -2, 50)) / 10.0,
            float(np.log10(mxm + 1)), float(outl))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--imatrix", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--ssim-table", default=None)
    ap.add_argument("--n-layers", type=int, default=0,
                    help="0 = estimate from tensor names")
    ap.add_argument("--labels", default=None,
                    help="jsonl damage file: unit tensors the imatrix does "
                         "not cover (MoE/SSM/MLA) still get feature rows")
    args = ap.parse_args()

    model = read_model(args.model)
    im = read_imatrix(args.imatrix)
    table = build_importance_table(im, model)
    ssim = load_ssim(args.ssim_table)

    # Ling scar (2026-10-25): the external imatrix covered 332/526 tensors —
    # no ssm_a/ssm_conv1d/ssm_dt entries — so 85 labelled units would have
    # had zero feature rows and the pool loader would refuse the pool.
    # Importance stays 0 (no imatrix evidence); n_elements and the weight
    # stats are real. No --labels = byte-identical old behaviour.
    if args.labels:
        lab_t = set()
        with open(args.labels) as f:
            for line in f:
                d = json.loads(line)
                if d.get("unit") == "BASELINE":
                    continue
                lab_t.update(d.get("tensors", []))
        missing = sorted(lab_t - set(table.keys()))
        for tname in missing:
            tinfo = model.get("tensors", {}).get(tname, {})
            table[tname] = {
                "importance_mean": 0.0, "importance_max": 0.0,
                "importance_min": 0.0, "importance_sum": 0.0,
                "n_elements": int(tinfo.get("n_elements", 0)),
                "type": tname,
            }
        print(f"  +{len(missing)} label tensors absent from the imatrix "
              f"→ zero-importance rows (table now {len(table)})", flush=True)

    reader = GGUFReader(args.model)
    wmap = {t.name: t for t in reader.tensors}

    names = sorted(table.keys())
    if args.n_layers <= 0:
        layers = [int(n.split(".")[1]) for n in names
                  if n.startswith("blk.") and n.split(".")[1].isdigit()]
        n_layers = max(layers) + 1 if layers else 28
    else:
        n_layers = args.n_layers

    X, kept = [], []
    for idx, n in enumerate(names):
        info = table[n]
        frag = 1.0 - ssim.get(n, 1.0)
        tm0, tx0 = info.get("importance_mean", 0.0), info.get("importance_max", 0.0)
        tmn = info.get("importance_min", 0.0)
        tsm = info.get("importance_sum", 0.0)
        econc = tx0 / (tm0 + 1e-12)  # peak/mean concentration ratio
        ws = None
        if n in wmap:
            try:
                v = _w_view(wmap[n])
                if v is not None and v.size:
                    ws = w_stats_stream(v)
            except Exception:
                ws = None
        row = featurize(n, info.get("n_elements", 0), tm0, tx0,
                        frag=frag, econc=econc, n_layers=n_layers, wstats=ws,
                        imp_min=tmn, imp_sum=tsm)
        X.append(row)
        kept.append(n)
        if (idx + 1) % 32 == 0 or idx + 1 == len(names):
            print(f"  [{idx + 1}/{len(names)}] {n}", flush=True)
    X = np.array(X, dtype=np.float64)
    tm = np.array([table[n].get("importance_mean", 0.0) for n in kept])
    logn = np.log10(np.array([table[n].get("n_elements", 1)
                              for n in kept]) + 1)
    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    np.savez(args.out, X=X, names=np.array(kept),
             timp_mean=np.array([tm.mean()]), timp_std=np.array([tm.std()]),
             logN_med=np.array([np.median(logn)]))
    print(f"{args.out}: {X.shape[0]} tensors x {X.shape[1]}, "
          f"n_layers={n_layers}", flush=True)


if __name__ == "__main__":
    main()