inferencelayer 0.2.9

//! GGUF WRITING + the native quantizer family — artifact PRODUCTION with zero external tools.
//!
//! This is the encode half of `gguf.rs`: pure-CPU Rust (no GPU, no llama.cpp binary, no Python)
//! that quantizes f32 weights into llama.cpp block formats and writes standard GGUF v3 files the
//! whole ecosystem (and our own `load_gguf_native`) can serve. The decode half is oracle-verified
//! bit-for-bit against llama.cpp's reference reader, so every encoder here is gated by
//! round-tripping through THOSE decoders — and the artifact-level judge is the KLD ruler, which
//! compares our encoder's artifacts against `llama-quantize`'s on equal terms (dev-gated test;
//! the runtime path has no such dependency).
//!
//! Encoder quality: the K-quant encoders fit per-sub-block (scale, min) by a few rounds of
//! quantize→least-squares refit (the same shape as llama.cpp's `make_qkx2_quants` iteration).
//! They are NOT bit-identical to llama.cpp's encoder — they don't need to be: the contract is
//! equal-or-better end-to-end fidelity at the same byte cost, measured, not byte equality.

use anyhow::{Context, Result, ensure};
use std::io::Write;

// ─── scale fitting ──────────────────────────────────────────────────────────────────────────

/// Fit `w ≈ s·q − m` with `q ∈ [0, qmax]`, `s > 0`, `m ≥ 0` (the K-quant affine form), the
/// `make_qkx2_quants` shape: importance-WEIGHTED (σ² + x², llama.cpp's default when no imatrix)
/// candidate-scale grid around the range solution, each candidate refined by weighted least
/// squares, best weighted-SSE wins. Returns (s, m).
fn fit_affine(w: &[f32], qmax: i32, im: Option<&[f32]>) -> (f32, f32) {
    let (mut lo, mut hi) = (f32::MAX, f32::MIN);
    let mut sum2 = 0f32;
    for &v in w {
        lo = lo.min(v);
        hi = hi.max(v);
        sum2 += v * v;
    }
    if hi == lo {
        return (0.0, (-lo).max(0.0));
    }
    let sigma2 = sum2 / w.len() as f32;
    let mut wt: Vec<f32> = w.iter().map(|&v| sigma2 + v * v).collect();
    if let Some(m) = im {
        for (t, &f) in wt.iter_mut().zip(m) {
            *t *= f;
        }
    }
    let m0 = (-lo).max(0.0);
    let s0 = (hi + m0) / qmax as f32;
    let sse = |s: f32, m: f32| -> f64 {
        let inv = if s != 0.0 { 1.0 / s } else { 0.0 };
        w.iter()
            .zip(&wt)
            .map(|(&v, &ww)| {
                let q = (((v + m) * inv).round()).clamp(0.0, qmax as f32);
                let r = (v - (s * q - m)) as f64;
                ww as f64 * r * r
            })
            .sum()
    };
    let refine = |mut s: f32, mut m: f32| -> (f32, f32) {
        for _ in 0..3 {
            let inv = if s != 0.0 { 1.0 / s } else { 0.0 };
            // weighted normal equations for w = s·q − m
            let (mut a11, mut a12, mut a22, mut b1, mut b2) = (0f64, 0f64, 0f64, 0f64, 0f64);
            for (&v, &ww) in w.iter().zip(&wt) {
                let q = ((((v + m) * inv).round()).clamp(0.0, qmax as f32)) as f64;
                let ww = ww as f64;
                a11 += ww * q * q;
                a12 -= ww * q;
                a22 += ww;
                b1 += ww * q * v as f64;
                b2 -= ww * v as f64;
            }
            let det = a11 * a22 - a12 * a12;
            if det.abs() < 1e-20 {
                break;
            }
            let ns = ((b1 * a22 - b2 * a12) / det) as f32;
            let nm = ((a11 * b2 - a12 * b1) / det) as f32;
            if ns <= 0.0 {
                break;
            }
            s = ns;
            m = nm.max(0.0);
        }
        (s.max(0.0), m.max(0.0))
    };
    let (mut best_s, mut best_m) = refine(s0, m0);
    let mut best = sse(best_s, best_m);
    // candidate grid: perturb the range scale like make_qkx2's is-loop, refine each
    let rdelta = s0 * 0.1;
    for is in -4i32..=4 {
        if is == 0 {
            continue;
        }
        let (cs, cm) = refine(s0 + is as f32 * rdelta, m0);
        let e = sse(cs, cm);
        if e < best {
            best = e;
            best_s = cs;
            best_m = cm;
        }
    }
    (best_s, best_m)
}

/// Fit `w ≈ s·q` with signed `q ∈ [−qmax−1, qmax]` (Q6_K's linear form): amax init + refit.
fn fit_linear(w: &[f32], qmax: i32, im: Option<&[f32]>) -> f32 {
    let amax = w.iter().fold(0f32, |a, &v| a.max(v.abs()));
    if amax == 0.0 {
        return 0.0;
    }
    let mut sum2 = 0f32;
    for &v in w {
        sum2 += v * v;
    }
    let sigma2 = sum2 / w.len() as f32;
    let mut s = amax / qmax as f32;
    for _ in 0..3 {
        let inv = 1.0 / s;
        let (mut num, mut den) = (0f64, 0f64);
        for (i, &v) in w.iter().enumerate() {
            let ww = ((sigma2 + v * v) * im.map_or(1.0, |m| m[i])) as f64;
            let q = ((v * inv).round().clamp(-(qmax as f32) - 1.0, qmax as f32)) as f64;
            num += ww * q * v as f64;
            den += ww * q * q;
        }
        if den <= 0.0 {
            break;
        }
        let ns = (num / den) as f32;
        if ns <= 0.0 {
            break;
        }
        s = ns;
    }
    s
}

fn f16b(v: f32) -> [u8; 2] {
    half::f16::from_f32(v).to_le_bytes()
}

// ─── legacy 32-block encoders (llama.cpp disk layouts) ──────────────────────────────────────

/// Imatrix-weighted scale search for the legacy signed formats (`make_qx_quants`' shape):
/// perturb the signed-amax scale, score by importance-weighted SSE (f16-rounded, matching
/// what disk will hold), keep the best. `lo`/`hi` are the clamp magnitudes (8/7 for Q4_0).
fn legacy_search(blk: &[f32], d0: f32, lo: f32, hi: f32, im: &[f32]) -> f32 {
    if d0 == 0.0 {
        return 0.0;
    }
    let sse = |d: f32| -> f64 {
        let df = half::f16::from_f32(d).to_f32();
        if df == 0.0 {
            return f64::MAX;
        }
        let inv = 1.0 / df;
        blk.iter()
            .zip(im)
            .map(|(&v, &ww)| {
                let q = (v * inv).round().clamp(-lo, hi);
                let r = (v - q * df) as f64;
                ww as f64 * r * r
            })
            .sum()
    };
    let (mut best_d, mut best) = (d0, sse(d0));
    for is in -10i32..=10 {
        if is == 0 {
            continue;
        }
        let c = d0 * (1.0 + 0.02 * is as f32);
        let e = sse(c);
        if e < best {
            best = e;
            best_d = c;
        }
    }
    best_d
}

/// Q8_0: 34-byte blocks — f16 `d = amax/127` + 32 `i8`.
pub fn enc_q8_0(w: &[f32]) -> Vec<u8> {
    assert!(w.len() % 32 == 0);
    let mut out = Vec::with_capacity(w.len() / 32 * 34);
    for blk in w.chunks_exact(32) {
        let amax = blk.iter().fold(0f32, |a, &v| a.max(v.abs()));
        let d = amax / 127.0;
        let db = half::f16::from_f32(d);
        let df = db.to_f32();
        let inv = if df != 0.0 { 1.0 / df } else { 0.0 };
        out.extend_from_slice(&db.to_le_bytes());
        for &v in blk {
            out.push(((v * inv).round().clamp(-127.0, 127.0)) as i8 as u8);
        }
    }
    out
}

/// Q4_0: 18-byte blocks — llama.cpp's signed-max convention (`d = max/−8`), byte `l` packs
/// elements `l` (low nibble) and `l+16` (high).
pub fn enc_q4_0(w: &[f32], cols: usize, im: Option<&[f32]>) -> Vec<u8> {
    assert!(w.len() % 32 == 0);
    let mut out = Vec::with_capacity(w.len() / 32 * 18);
    for (k, blk) in w.chunks_exact(32).enumerate() {
        let (mut amax, mut smax) = (0f32, 0f32);
        for &v in blk {
            if v.abs() > amax {
                amax = v.abs();
                smax = v;
            }
        }
        let d0 = smax / -8.0;
        let d = match im {
            Some(m) => legacy_search(blk, d0, 8.0, 7.0, &m[(k * 32) % cols..(k * 32) % cols + 32]),
            None => d0, // no imatrix: llama.cpp's exact signed-amax convention
        };
        let db = half::f16::from_f32(d);
        let df = db.to_f32();
        let inv = if df != 0.0 { 1.0 / df } else { 0.0 };
        out.extend_from_slice(&db.to_le_bytes());
        for l in 0..16 {
            let q0 = ((blk[l] * inv).round() as i32 + 8).clamp(0, 15) as u8;
            let q1 = ((blk[l + 16] * inv).round() as i32 + 8).clamp(0, 15) as u8;
            out.push(q0 | (q1 << 4));
        }
    }
    out
}

/// Q5_0: 22-byte blocks — `d = max/−16`, 5th bits packed into a u32.
pub fn enc_q5_0(w: &[f32], cols: usize, im: Option<&[f32]>) -> Vec<u8> {
    assert!(w.len() % 32 == 0);
    let mut out = Vec::with_capacity(w.len() / 32 * 22);
    for (k, blk) in w.chunks_exact(32).enumerate() {
        let (mut amax, mut smax) = (0f32, 0f32);
        for &v in blk {
            if v.abs() > amax {
                amax = v.abs();
                smax = v;
            }
        }
        let d0 = smax / -16.0;
        let d = match im {
            Some(m) => {
                legacy_search(blk, d0, 16.0, 15.0, &m[(k * 32) % cols..(k * 32) % cols + 32])
            }
            None => d0, // no imatrix: llama.cpp's exact signed-amax convention
        };
        let db = half::f16::from_f32(d);
        let df = db.to_f32();
        let inv = if df != 0.0 { 1.0 / df } else { 0.0 };
        let q: Vec<u8> = blk
            .iter()
            .map(|&v| ((v * inv).round() as i32 + 16).clamp(0, 31) as u8)
            .collect();
        let mut qh = 0u32;
        for l in 0..16 {
            qh |= (((q[l] >> 4) & 1) as u32) << l;
            qh |= (((q[l + 16] >> 4) & 1) as u32) << (l + 16);
        }
        out.extend_from_slice(&db.to_le_bytes());
        out.extend_from_slice(&qh.to_le_bytes());
        for l in 0..16 {
            out.push((q[l] & 0x0F) | ((q[l + 16] & 0x0F) << 4));
        }
    }
    out
}

// ─── IQ4 nonlinear-codebook encoders ────────────────────────────────────────────────────────

/// Nearest index in the 16-entry IQ4 codebook (kvalues ascending; 16-wide scan).
fn iq4_idx(v: f32) -> usize {
    let mut bi = 0;
    let mut bd = f32::MAX;
    for (i, &k) in crate::iq_tables::IQ4_KVALUES.iter().enumerate() {
        let d = (v - k as f32).abs();
        if d < bd {
            bd = d;
            bi = i;
        }
    }
    bi
}

/// Weighted LS scale fit through the codebook (`quantize_row_iq4_nl`'s shape): init from the
/// signed max against kvalues[0] = −127, iterate index-assign → refit, over a candidate grid.
/// Importance defaults to σ²+x² like the K-quant fits.
fn fit_iq4(blk: &[f32], im: Option<&[f32]>) -> f32 {
    let (mut amax, mut smax) = (0f32, 0f32);
    let mut sum2 = 0f32;
    for &v in blk {
        if v.abs() > amax {
            amax = v.abs();
            smax = v;
        }
        sum2 += v * v;
    }
    if amax == 0.0 {
        return 0.0;
    }
    let sigma2 = sum2 / blk.len() as f32;
    let wt = |i: usize, v: f32| -> f64 {
        ((sigma2 + v * v) * im.map_or(1.0, |m| m[i])) as f64
    };
    let d0 = smax / crate::iq_tables::IQ4_KVALUES[0] as f32;
    let sse = |d: f32| -> f64 {
        let df = half::f16::from_f32(d).to_f32();
        if df == 0.0 {
            return f64::MAX;
        }
        let inv = 1.0 / df;
        blk.iter()
            .enumerate()
            .map(|(i, &v)| {
                let k = crate::iq_tables::IQ4_KVALUES[iq4_idx(v * inv)] as f32;
                let r = (v - df * k) as f64;
                wt(i, v) * r * r
            })
            .sum()
    };
    let refine = |mut d: f32| -> f32 {
        for _ in 0..3 {
            if d == 0.0 {
                break;
            }
            let inv = 1.0 / d;
            let (mut num, mut den) = (0f64, 0f64);
            for (i, &v) in blk.iter().enumerate() {
                let k = crate::iq_tables::IQ4_KVALUES[iq4_idx(v * inv)] as f32 as f64;
                let w = wt(i, v);
                num += w * k * v as f64;
                den += w * k * k;
            }
            if den <= 0.0 {
                break;
            }
            let nd = (num / den) as f32;
            if nd == 0.0 {
                break;
            }
            d = nd;
        }
        d
    };
    let (mut best_d, mut best) = (refine(d0), 0f64);
    best = sse(best_d);
    for is in -4i32..=4 {
        if is == 0 {
            continue;
        }
        let c = refine(d0 * (1.0 + 0.05 * is as f32));
        let e = sse(c);
        if e < best {
            best = e;
            best_d = c;
        }
    }
    best_d
}

/// IQ4_NL: 18-byte/32 blocks — f16 d + LUT nibbles in Q4_0's byte layout.
pub fn enc_iq4_nl(w: &[f32], cols: usize, im: Option<&[f32]>) -> Vec<u8> {
    assert!(w.len() % 32 == 0);
    let mut out = Vec::with_capacity(w.len() / 32 * 18);
    for (k, blk) in w.chunks_exact(32).enumerate() {
        let off = (k * 32) % cols;
        let d = fit_iq4(blk, im.map(|m| &m[off..off + 32]));
        let db = half::f16::from_f32(d);
        let df = db.to_f32();
        let inv = if df != 0.0 { 1.0 / df } else { 0.0 };
        out.extend_from_slice(&db.to_le_bytes());
        for l in 0..16 {
            let q0 = iq4_idx(blk[l] * inv) as u8;
            let q1 = iq4_idx(blk[l + 16] * inv) as u8;
            out.push(q0 | (q1 << 4));
        }
    }
    out
}

/// IQ4_XS: 136-byte/256 superblocks — 8 sub-block LUT fits super-quantized to 6-bit
/// (−32-biased) scales under one f16 d; elements re-indexed against the DECODED effective
/// scale (the K-quant discipline).
pub fn enc_iq4_xs(w: &[f32], cols: usize, im: Option<&[f32]>) -> Vec<u8> {
    assert!(w.len() % 256 == 0);
    let mut out = Vec::with_capacity(w.len() / 256 * 136);
    for (kb, sb) in w.chunks_exact(256).enumerate() {
        let off = (kb * 256) % cols;
        let mut ss = [0f32; 8];
        for b in 0..8 {
            ss[b] = fit_iq4(
                &sb[b * 32..(b + 1) * 32],
                im.map(|m| &m[off + b * 32..off + (b + 1) * 32]),
            );
        }
        let amax = ss.iter().fold(0f32, |a, &v| a.max(v.abs()));
        let d = amax / 31.0;
        let db = half::f16::from_f32(d);
        let df = db.to_f32();
        let mut ls = [32u8; 8];
        for b in 0..8 {
            ls[b] = if df > 0.0 {
                (((ss[b] / df).round() as i32).clamp(-32, 31) + 32) as u8
            } else {
                32
            };
        }
        out.extend_from_slice(&db.to_le_bytes());
        let mut sh = 0u16;
        for (b, &l) in ls.iter().enumerate() {
            sh |= (((l >> 4) & 3) as u16) << (2 * b);
        }
        out.extend_from_slice(&sh.to_le_bytes());
        for b in 0..4 {
            out.push((ls[2 * b] & 0x0F) | ((ls[2 * b + 1] & 0x0F) << 4));
        }
        for b in 0..8 {
            let eff = df * (ls[b] as i8 as i32 - 32) as f32;
            let inv = if eff != 0.0 { 1.0 / eff } else { 0.0 };
            let blk = &sb[b * 32..(b + 1) * 32];
            for l in 0..16 {
                let q0 = iq4_idx(blk[l] * inv) as u8;
                let q1 = iq4_idx(blk[l + 16] * inv) as u8;
                out.push(q0 | (q1 << 4));
            }
        }
    }
    out
}

// ─── K-quant superblock encoders ────────────────────────────────────────────────────────────

/// Pack eight 6-bit (sc, mn) pairs into the 12-byte `get_scale_min` layout (its exact inverse).
fn pack_scale_min(sc: &[u8; 8], mn: &[u8; 8]) -> [u8; 12] {
    let mut b = [0u8; 12];
    for i in 0..4 {
        b[i] = (sc[i] & 63) | ((sc[i + 4] >> 4) << 6);
        b[i + 4] = (mn[i] & 63) | ((mn[i + 4] >> 4) << 6);
        b[i + 8] = (sc[i + 4] & 0x0F) | ((mn[i + 4] & 0x0F) << 4);
    }
    b
}

/// Shared Q4_K/Q5_K encode: fit 8 affine sub-blocks, super-quantize the scales to 6 bit,
/// re-quantize elements against the DECODED effective scales (so encode matches what the
/// oracle-verified decoder will reconstruct). Returns per-superblock (d, dmin, sc, mn, q).
///
/// Mutation-tested (2026-07): layout/sign/shift mutations are all CAUGHT by the lib gates;
/// the super-scale divisors (63, 127) are llama.cpp's convention but neighboring divisors
/// are EQUIVALENT MUTANTS — the elements re-quantize against the decoded effective scales,
/// so any d keeping sc in 6-bit range is a valid encoding (measured RMSE differs at the 5th
/// decimal with varying sign). Deliberately not pinned.
fn kquant_fit(
    sb: &[f32],
    qmax: i32,
    im: Option<&[f32]>,
) -> ([u8; 2], [u8; 2], [u8; 8], [u8; 8], [u8; 256]) {
    let mut ss = [0f32; 8];
    let mut ms = [0f32; 8];
    for s in 0..8 {
        let (sv, mv) = fit_affine(
            &sb[s * 32..(s + 1) * 32],
            qmax,
            im.map(|m| &m[s * 32..(s + 1) * 32]),
        );
        ss[s] = sv;
        ms[s] = mv;
    }
    let smax = ss.iter().fold(0f32, |a, &v| a.max(v));
    let mmax = ms.iter().fold(0f32, |a, &v| a.max(v));
    let d = smax / 63.0;
    let dmin = mmax / 63.0;
    let (db, mb) = (half::f16::from_f32(d), half::f16::from_f32(dmin));
    let (df, mf) = (db.to_f32(), mb.to_f32());
    let mut sc = [0u8; 8];
    let mut mn = [0u8; 8];
    for s in 0..8 {
        sc[s] = if df > 0.0 {
            ((ss[s] / df).round() as i32).clamp(0, 63) as u8
        } else {
            0
        };
        mn[s] = if mf > 0.0 {
            ((ms[s] / mf).round() as i32).clamp(0, 63) as u8
        } else {
            0
        };
    }
    let mut q = [0u8; 256];
    for s in 0..8 {
        let eff_s = df * sc[s] as f32;
        let eff_m = mf * mn[s] as f32;
        let inv = if eff_s > 0.0 { 1.0 / eff_s } else { 0.0 };
        for l in 0..32 {
            let v = sb[s * 32 + l];
            q[s * 32 + l] = (((v + eff_m) * inv).round()).clamp(0.0, qmax as f32) as u8;
        }
    }
    (db.to_le_bytes(), mb.to_le_bytes(), sc, mn, q)
}

/// Q4_K: 144-byte superblocks.
pub fn enc_q4_k(w: &[f32], cols: usize, im: Option<&[f32]>) -> Vec<u8> {
    assert!(w.len() % 256 == 0);
    let mut out = Vec::with_capacity(w.len() / 256 * 144);
    for (k, sb) in w.chunks_exact(256).enumerate() {
        let off = (k * 256) % cols;
        let (db, mb, sc, mn, q) = kquant_fit(sb, 15, im.map(|m| &m[off..off + 256]));
        out.extend_from_slice(&db);
        out.extend_from_slice(&mb);
        out.extend_from_slice(&pack_scale_min(&sc, &mn));
        // nibble planes: sub-block s → byte group s/2, low nibble s even / high s odd
        for g in 0..4 {
            for l in 0..32 {
                out.push((q[g * 64 + l] & 0x0F) | ((q[g * 64 + 32 + l] & 0x0F) << 4));
            }
        }
    }
    out
}

/// Q5_K: 176-byte superblocks — Q4_K's fit at qmax 31 plus the qh bit plane.
pub fn enc_q5_k(w: &[f32], cols: usize, im: Option<&[f32]>) -> Vec<u8> {
    assert!(w.len() % 256 == 0);
    let mut out = Vec::with_capacity(w.len() / 256 * 176);
    for (k, sb) in w.chunks_exact(256).enumerate() {
        let off = (k * 256) % cols;
        let (db, mb, sc, mn, q) = kquant_fit(sb, 31, im.map(|m| &m[off..off + 256]));
        out.extend_from_slice(&db);
        out.extend_from_slice(&mb);
        out.extend_from_slice(&pack_scale_min(&sc, &mn));
        let mut qh = [0u8; 32];
        for s in 0..8 {
            for l in 0..32 {
                qh[l] |= ((q[s * 32 + l] >> 4) & 1) << s;
            }
        }
        out.extend_from_slice(&qh);
        for g in 0..4 {
            for l in 0..32 {
                out.push((q[g * 64 + l] & 0x0F) | ((q[g * 64 + 32 + l] & 0x0F) << 4));
            }
        }
    }
    out
}

// ─── grid-codebook encoders (IQ3_XXS / IQ2_XS) ──────────────────────────────────────────────
// Brute-force weighted nearest-row search over the (small) codebooks — llama.cpp reaches the
// same argmin through precomputed neighbour maps; at artifact-production time the exhaustive
// scan is simpler and parallelizes over superblocks. Signs are ksign-coded (7 bits per 8
// elems, 8th = parity): odd-parity sign patterns flip the least-importance element, exactly
// llama.cpp's move.

/// (ksign 7-bit index, index of the flipped element or usize::MAX). `neg[k]` = element k is
/// negative; flip cost is weighted |w| so the least valuable sign absorbs the parity error.
fn ksign_encode(neg: &mut [bool; 8], w: &[f32], wt: &[f32]) -> u32 {
    let par = neg.iter().filter(|&&b| b).count() % 2;
    if par == 1 {
        let mut fi = 0;
        let mut fc = f32::MAX;
        for k in 0..8 {
            let c = wt[k] * w[k].abs();
            if c < fc {
                fc = c;
                fi = k;
            }
        }
        neg[fi] = !neg[fi];
    }
    let mut idx = 0u32;
    for (k, &b) in neg.iter().take(7).enumerate() {
        idx |= (b as u32) << k;
    }
    idx
}

/// Weighted nearest grid row for a |target| vector; rows are `rl` positive i8 magnitudes.
fn best_grid_row(grid: &[i8], rl: usize, tgt: &[f32], wt: &[f32]) -> usize {
    let mut bi = 0usize;
    let mut be = f32::MAX;
    for (ri, row) in grid.chunks_exact(rl).enumerate() {
        let mut e = 0f32;
        for k in 0..rl {
            let d = tgt[k] - row[k] as f32;
            e += wt[k] * d * d;
        }
        if e < be {
            be = e;
            bi = ri;
        }
    }
    bi
}

fn iq_weights(sb: &[f32], im: Option<&[f32]>) -> Vec<f32> {
    let sigma2 = sb.iter().map(|v| v * v).sum::<f32>() / sb.len() as f32;
    sb.iter()
        .enumerate()
        .map(|(i, &v)| (sigma2 + v * v) * im.map_or(1.0, |m| m[i]))
        .collect()
}

/// Parallel-over-superblocks driver: `enc(sb, im_sb) -> [u8; BB]` applied on row-aligned
/// chunks via scoped threads (no rayon feature juggling).
fn par_encode_sb(
    w: &[f32],
    cols: usize,
    im: Option<&[f32]>,
    bb: usize,
    enc: impl Fn(&[f32], Option<&[f32]>) -> Vec<u8> + Sync,
) -> Vec<u8> {
    let nsb = w.len() / 256;
    let nthread = std::thread::available_parallelism().map_or(4, |n| n.get()).min(16);
    let chunk = nsb.div_ceil(nthread);
    let mut out = vec![0u8; nsb * bb];
    std::thread::scope(|sc| {
        for (t, dst) in out.chunks_mut(chunk * bb).enumerate() {
            let enc = &enc;
            sc.spawn(move || {
                for (j, d) in dst.chunks_mut(bb).enumerate() {
                    let k = t * chunk + j;
                    let off = (k * 256) % cols;
                    let e = enc(&w[k * 256..(k + 1) * 256], im.map(|m| &m[off..off + 256]));
                    d.copy_from_slice(&e);
                }
            });
        }
    });
    out
}

/// IQ3_XXS: 98-byte superblocks — 8 groups of 32, per group a 4-bit scale + 4 ksign words
/// covering 8 elems each; magnitudes from the 256x4 grid. Two-round scale iteration.
pub fn enc_iq3_xxs(w: &[f32], cols: usize, im: Option<&[f32]>) -> Vec<u8> {
    assert!(w.len() % 256 == 0);
    let grid = &crate::iq_tables::IQ3_XXS_GRID[..];
    par_encode_sb(w, cols, im, 98, |sb, imsb| {
        let wt = iq_weights(sb, imsb);
        // superblock d from the loudest group's ideal sub-scale
        let mut dbmax = 0f32;
        for g in 0..8 {
            let amax = sb[g * 32..(g + 1) * 32]
                .iter()
                .fold(0f32, |a, &v| a.max(v.abs()));
            dbmax = dbmax.max(amax / 62.0);
        }
        let mut d = dbmax / 7.75;
        let mut best: Option<(f32, Vec<u8>)> = None;
        for round in 0..3 {
            if d <= 0.0 {
                d = 1e-8;
            }
            let df = half::f16::from_f32(d).to_f32().max(1e-12);
            let mut blk = Vec::with_capacity(98);
            blk.extend_from_slice(&half::f16::from_f32(d).to_le_bytes());
            let mut qs = [0u8; 64];
            let mut sw = [0u32; 8];
            let mut sse = 0f64;
            let (mut num, mut den) = (0f64, 0f64);
            for g in 0..8 {
                let gsl = &sb[g * 32..(g + 1) * 32];
                let gwt = &wt[g * 32..(g + 1) * 32];
                let amax = gsl.iter().fold(0f32, |a, &v| a.max(v.abs()));
                let want = amax / 62.0;
                let s0 = ((want / (0.5 * df) - 0.5).round() as i32).clamp(0, 15);
                // per-group scale SEARCH: the amax pick loses ~10% KLD to llama.cpp's
                // SSE-searched scales on body tensors — try the neighbourhood, keep best.
                let mut gbest: Option<(f64, u32)> = None;
                for ds in -1i32..=1 {
                    let s4c = (s0 + ds).clamp(0, 15) as u32;
                    let dbc = df * (0.5 + s4c as f32) * 0.5;
                    let mut ge = 0f64;
                    for t8 in 0..8 {
                        let e0 = t8 * 4;
                        let mut tgt = [0f32; 4];
                        for k in 0..4 {
                            tgt[k] = gsl[e0 + k].abs() / dbc;
                        }
                        let row = best_grid_row(grid, 4, &tgt, &gwt[e0..e0 + 4]);
                        for k in 0..4 {
                            let r = dbc * grid[row * 4 + k] as f32;
                            let dd = gsl[e0 + k].abs() - r;
                            ge += gwt[e0 + k] as f64 * (dd as f64) * (dd as f64);
                        }
                    }
                    if gbest.is_none_or(|(b, _)| ge < b) {
                        gbest = Some((ge, s4c));
                    }
                }
                let s4 = gbest.unwrap().1;
                let db = df * (0.5 + s4 as f32) * 0.5;
                sw[g] = s4 << 28;
                for j in 0..4 {
                    let e0 = j * 8;
                    let mut neg = [false; 8];
                    for k in 0..8 {
                        neg[k] = gsl[e0 + k] < 0.0;
                    }
                    let sidx = ksign_encode(&mut neg, &gsl[e0..e0 + 8], &gwt[e0..e0 + 8]);
                    sw[g] |= sidx << (7 * j);
                    for t in 0..2 {
                        let mut tgt = [0f32; 4];
                        for k in 0..4 {
                            tgt[k] = gsl[e0 + t * 4 + k].abs() / db;
                        }
                        let row = best_grid_row(grid, 4, &tgt, &gwt[e0 + t * 4..e0 + t * 4 + 4]);
                        qs[g * 8 + j * 2 + t] = row as u8;
                        for k in 0..4 {
                            let sgnf = if neg[t * 4 + k] { -1.0f32 } else { 1.0 };
                            let r = (0.5 + s4 as f32) * 0.5 * grid[row * 4 + k] as f32 * sgnf;
                            let v = sb[g * 32 + e0 + t * 4 + k];
                            let ww = gwt[e0 + t * 4 + k] as f64;
                            sse += ww * ((v - df * r) as f64).powi(2);
                            num += ww * r as f64 * v as f64;
                            den += ww * (r as f64) * (r as f64);
                        }
                    }
                }
            }
            blk.extend_from_slice(bytemuck::cast_slice(&qs));
            for v in sw {
                blk.extend_from_slice(&v.to_le_bytes());
            }
            if best.as_ref().is_none_or(|(b, _)| (sse as f32) < *b) {
                best = Some((sse as f32, blk));
            }
            if round < 2 && den > 0.0 {
                d = (num / den) as f32;
            }
        }
        best.unwrap().1
    })
}

/// IQ2_XS: 74-byte superblocks — 32 u16 words (9-bit row into the 512x8 grid + 7-bit ksign),
/// 16 4-bit scales (one per 2 words). Same two-round scale iteration.
pub fn enc_iq2_xs(w: &[f32], cols: usize, im: Option<&[f32]>) -> Vec<u8> {
    assert!(w.len() % 256 == 0);
    let grid = &crate::iq_tables::IQ2_XS_GRID[..];
    par_encode_sb(w, cols, im, 74, |sb, imsb| {
        let wt = iq_weights(sb, imsb);
        let mut dbmax = 0f32;
        for s16 in 0..16 {
            let amax = sb[s16 * 16..(s16 + 1) * 16]
                .iter()
                .fold(0f32, |a, &v| a.max(v.abs()));
            dbmax = dbmax.max(amax / 43.0);
        }
        let mut d = dbmax / 3.875; // eff max = d*(0.5+15)*0.25
        let mut best: Option<(f32, Vec<u8>)> = None;
        for round in 0..3 {
            if d <= 0.0 {
                d = 1e-8;
            }
            let df = half::f16::from_f32(d).to_f32().max(1e-12);
            let mut words = [0u16; 32];
            let mut scales = [0u8; 8];
            let mut sse = 0f64;
            let (mut num, mut den) = (0f64, 0f64);
            for s16 in 0..16 {
                let gsl = &sb[s16 * 16..(s16 + 1) * 16];
                let gwt = &wt[s16 * 16..(s16 + 1) * 16];
                let amax = gsl.iter().fold(0f32, |a, &v| a.max(v.abs()));
                let want = amax / 43.0;
                let s0 = ((want / (0.25 * df) - 0.5).round() as i32).clamp(0, 15);
                // per-group scale SEARCH (the IQ3 lesson: amax-only picks cost double-digit
                // KLD at low bits; SSE over the candidate neighbourhood decides).
                // Joint row<->scale refinement (the −19% IQ2 lever): fan scale candidates,
                // pick rows per candidate; then the EXACT LS scale over the chosen rows,
                // re-rounded and re-searched once if it moves.
                let ge_at = |s4c: u32| -> f64 {
                    let dbc = df * (0.5 + s4c as f32) * 0.25;
                    let mut ge = 0f64;
                    for h in 0..2 {
                        let e0 = h * 8;
                        let mut tgt = [0f32; 8];
                        for k in 0..8 {
                            tgt[k] = gsl[e0 + k].abs() / dbc;
                        }
                        let row = best_grid_row(grid, 8, &tgt, &gwt[e0..e0 + 8]);
                        for k in 0..8 {
                            let r = dbc * grid[row * 8 + k] as f32;
                            let dd = gsl[e0 + k].abs() - r;
                            ge += gwt[e0 + k] as f64 * (dd as f64) * (dd as f64);
                        }
                    }
                    ge
                };
                let mut gbest: Option<(f64, u32)> = None;
                for ds in -2i32..=2 {
                    let s4c = (s0 + ds).clamp(0, 15) as u32;
                    if gbest.as_ref().is_some_and(|&(_, b)| b == s4c) {
                        continue;
                    }
                    let ge = ge_at(s4c);
                    if gbest.is_none_or(|(b, _)| ge < b) {
                        gbest = Some((ge, s4c));
                    }
                }
                // exact LS scale over the winner's rows, re-rounded
                let (mut gbe, mut s4) = gbest.unwrap();
                {
                    let dbc = df * (0.5 + s4 as f32) * 0.25;
                    let (mut gn, mut gd) = (0f64, 0f64);
                    for h in 0..2 {
                        let e0 = h * 8;
                        let mut tgt = [0f32; 8];
                        for k in 0..8 {
                            tgt[k] = gsl[e0 + k].abs() / dbc;
                        }
                        let row = best_grid_row(grid, 8, &tgt, &gwt[e0..e0 + 8]);
                        for k in 0..8 {
                            let gv = grid[row * 8 + k] as f64;
                            let ww = gwt[e0 + k] as f64;
                            gn += ww * gv * gsl[e0 + k].abs() as f64;
                            gd += ww * gv * gv;
                        }
                    }
                    if gd > 0.0 {
                        let dbx = (gn / gd) as f32;
                        let s4x = ((dbx / (0.25 * df) - 0.5).round() as i32).clamp(0, 15) as u32;
                        if s4x != s4 {
                            let gex = ge_at(s4x);
                            if gex < gbe {
                                gbe = gex;
                                s4 = s4x;
                            }
                        }
                    }
                    let _ = gbe;
                }
                let db = df * (0.5 + s4 as f32) * 0.25;
                scales[s16 / 2] |= (s4 as u8) << (4 * (s16 % 2));
                for h in 0..2 {
                    let e0 = h * 8;
                    let mut neg = [false; 8];
                    for k in 0..8 {
                        neg[k] = gsl[e0 + k] < 0.0;
                    }
                    let sidx = ksign_encode(&mut neg, &gsl[e0..e0 + 8], &gwt[e0..e0 + 8]);
                    let mut tgt = [0f32; 8];
                    for k in 0..8 {
                        tgt[k] = gsl[e0 + k].abs() / db;
                    }
                    let row = best_grid_row(grid, 8, &tgt, &gwt[e0..e0 + 8]);
                    words[s16 * 2 + h] = (row as u16) | ((sidx as u16) << 9);
                    for k in 0..8 {
                        let sgnf = if neg[k] { -1.0f32 } else { 1.0 };
                        let r = (0.5 + s4 as f32) * 0.25 * grid[row * 8 + k] as f32 * sgnf;
                        let v = gsl[e0 + k];
                        let ww = gwt[e0 + k] as f64;
                        sse += ww * ((v - df * r) as f64).powi(2);
                        num += ww * r as f64 * v as f64;
                        den += ww * (r as f64) * (r as f64);
                    }
                }
            }
            let mut blk = Vec::with_capacity(74);
            blk.extend_from_slice(&half::f16::from_f32(d).to_le_bytes());
            for v in words {
                blk.extend_from_slice(&v.to_le_bytes());
            }
            blk.extend_from_slice(&scales);
            if best.as_ref().is_none_or(|(b, _)| (sse as f32) < *b) {
                best = Some((sse as f32, blk));
            }
            if round < 2 && den > 0.0 {
                d = (num / den) as f32;
            }
        }
        best.unwrap().1
    })
}

/// IQ3_S: 110-byte superblocks — 512x4 grid (9th row bit in qh), EXPLICIT sign bytes (no
/// parity code), 8 4-bit scales with db = d*(1+2s). The easiest grid encoder: signs are free.
pub fn enc_iq3_s(w: &[f32], cols: usize, im: Option<&[f32]>) -> Vec<u8> {
    assert!(w.len() % 256 == 0);
    let grid = &crate::iq_tables::IQ3_S_GRID[..];
    par_encode_sb(w, cols, im, 110, |sb, imsb| {
        let wt = iq_weights(sb, imsb);
        let mut dbmax = 0f32;
        for g in 0..8 {
            let amax = sb[g * 32..(g + 1) * 32]
                .iter()
                .fold(0f32, |a, &v| a.max(v.abs()));
            dbmax = dbmax.max(amax / 15.0);
        }
        let mut d = dbmax / 31.0; // db = d*(1+2s), s<=15 -> max 31d
        let mut best: Option<(f32, Vec<u8>)> = None;
        for round in 0..3 {
            if d <= 0.0 {
                d = 1e-8;
            }
            let df = half::f16::from_f32(d).to_f32().max(1e-12);
            let mut qs = [0u8; 64];
            let mut qh = [0u8; 8];
            let mut sg = [0u8; 32];
            let mut sc = [0u8; 4];
            let mut sse = 0f64;
            let (mut num, mut den) = (0f64, 0f64);
            for g in 0..8 {
                let gsl = &sb[g * 32..(g + 1) * 32];
                let gwt = &wt[g * 32..(g + 1) * 32];
                let amax = gsl.iter().fold(0f32, |a, &v| a.max(v.abs()));
                let s4 = (((amax / 15.0 / df - 1.0) / 2.0).round() as i32).clamp(0, 15) as u32;
                let db = df * (1 + 2 * s4) as f32;
                sc[g / 2] |= (s4 as u8) << (4 * (g % 2));
                for t in 0..8 {
                    let e0 = t * 4;
                    let mut tgt = [0f32; 4];
                    for k in 0..4 {
                        tgt[k] = gsl[e0 + k].abs() / db;
                        if gsl[e0 + k] < 0.0 {
                            sg[g * 4 + t / 2] |= 1 << ((t % 2) * 4 + k);
                        }
                    }
                    let row = best_grid_row(grid, 4, &tgt, &gwt[e0..e0 + 4]);
                    let qi = g * 8 + t;
                    qs[qi] = (row & 0xFF) as u8;
                    if row >= 256 {
                        qh[qi / 8] |= 1 << (qi % 8);
                    }
                    for k in 0..4 {
                        let sgnf = if gsl[e0 + k] < 0.0 { -1.0f32 } else { 1.0 };
                        let r = (1 + 2 * s4) as f32 * grid[row * 4 + k] as f32 * sgnf;
                        let v = gsl[e0 + k];
                        let ww = gwt[e0 + k] as f64;
                        sse += ww * ((v - df * r) as f64).powi(2);
                        num += ww * r as f64 * v as f64;
                        den += ww * (r as f64) * (r as f64);
                    }
                }
            }
            let mut blk = Vec::with_capacity(110);
            blk.extend_from_slice(&half::f16::from_f32(d).to_le_bytes());
            blk.extend_from_slice(&qs);
            blk.extend_from_slice(&qh);
            blk.extend_from_slice(&sg);
            blk.extend_from_slice(&sc);
            if best.as_ref().is_none_or(|(b, _)| (sse as f32) < *b) {
                best = Some((sse as f32, blk));
            }
            if round < 2 && den > 0.0 {
                d = (num / den) as f32;
            }
        }
        best.unwrap().1
    })
}

/// Q6_K: 210-byte superblocks/// Q6_K: 210-byte superblocks/// Q6_K: 210-byte superblocks — sixteen signed-linear 16-element groups, i8 group scales,
/// f16 super scale.
pub fn enc_q6_k(w: &[f32], cols: usize, im: Option<&[f32]>) -> Vec<u8> {
    assert!(w.len() % 256 == 0);
    let mut out = Vec::with_capacity(w.len() / 256 * 210);
    for (k, sb) in w.chunks_exact(256).enumerate() {
        let off = (k * 256) % cols;
        let mut gs = [0f32; 16];
        for g in 0..16 {
            gs[g] = fit_linear(
                &sb[g * 16..(g + 1) * 16],
                31,
                im.map(|m| &m[off + g * 16..off + (g + 1) * 16]),
            );
        }
        let smax = gs.iter().fold(0f32, |a, &v| a.max(v.abs()));
        let d = smax / 127.0;
        let db = half::f16::from_f32(d);
        let df = db.to_f32();
        let mut sc = [0i8; 16];
        for g in 0..16 {
            sc[g] = if df > 0.0 {
                ((gs[g] / df).round() as i32).clamp(-128, 127) as i8
            } else {
                0
            };
        }
        // quantize against the DECODED effective scale, element order per the decoder:
        // e = half·128 + quarter·32 + l ; q6 = clamp(round(w/(d·sc)), −32..31) + 32
        let mut q6 = [0u8; 256];
        for e in 0..256 {
            let g = e / 16;
            let eff = df * sc[g] as f32;
            let q = if eff != 0.0 {
                ((sb[e] / eff).round()).clamp(-32.0, 31.0) as i32 + 32
            } else {
                32
            };
            q6[e] = q as u8;
        }
        let mut ql = [0u8; 128];
        let mut qh = [0u8; 64];
        for e in 0..256 {
            let (half_i, r) = (e / 128, e % 128);
            let (quarter, l) = (r / 32, r % 32);
            let lo = q6[e] & 0x0F;
            let hi = (q6[e] >> 4) & 0x03;
            let qli = half_i * 64 + (quarter % 2) * 32 + l;
            ql[qli] |= lo << (4 * (quarter / 2));
            qh[half_i * 32 + l] |= hi << (2 * quarter);
        }
        out.extend_from_slice(&ql);
        out.extend_from_slice(&qh);
        for g in 0..16 {
            out.push(sc[g] as u8);
        }
        out.extend_from_slice(&db.to_le_bytes());
    }
    out
}

/// F16 tensor bytes.
pub fn enc_f16(w: &[f32]) -> Vec<u8> {
    let mut out = Vec::with_capacity(w.len() * 2);
    for &v in w {
        out.extend_from_slice(&f16b(v));
    }
    out
}

/// Encode by ggml type id (the ids `gguf.rs` decodes). `cols` is the row length (block-count
/// validity is per ROW — blocks never straddle rows).
pub fn encode_ggml(w: &[f32], cols: usize, ty: u32) -> Result<Vec<u8>> {
    encode_ggml_im(w, cols, ty, None)
}

/// [`encode_ggml`] with an optional per-COLUMN importance row (an imatrix diagonal, len
/// `cols`): K-quant fits weight their SSE by it; legacy Q4_0/Q5_0 run a weighted scale
/// search (their no-imatrix path stays the bit-exact llama.cpp amax convention). Q8_0 and
/// f16 ignore importance (llama.cpp does the same — nothing to trade at 8+ bits).
pub fn encode_ggml_im(w: &[f32], cols: usize, ty: u32, im: Option<&[f32]>) -> Result<Vec<u8>> {
    let be = match ty {
        1 => 1,
        2 | 6 | 8 | 20 => 32,
        12 | 13 | 14 | 17 | 18 | 21 | 23 => 256,
        other => anyhow::bail!("encode_ggml: unsupported ggml type {other}"),
    };
    ensure!(cols % be == 0, "cols {cols} not a multiple of block {be} for type {ty}");
    if let Some(m) = im {
        ensure!(m.len() == cols, "imatrix len {} != cols {cols}", m.len());
    }
    Ok(match ty {
        1 => enc_f16(w),
        2 => enc_q4_0(w, cols, im),
        6 => enc_q5_0(w, cols, im),
        8 => enc_q8_0(w),
        12 => enc_q4_k(w, cols, im),
        13 => enc_q5_k(w, cols, im),
        14 => enc_q6_k(w, cols, im),
        17 => enc_iq2_xs(w, cols, im),
        18 => enc_iq3_xxs(w, cols, im),
        21 => enc_iq3_s(w, cols, im),
        20 => enc_iq4_nl(w, cols, im),
        23 => enc_iq4_xs(w, cols, im),
        _ => unreachable!(),
    })
}

/// Bytes per block / elements per block for a ggml type id.
pub fn ggml_block(ty: u32) -> (usize, usize) {
    match ty {
        1 => (1, 2),
        2 => (32, 18),
        6 => (32, 22),
        8 => (32, 34),
        12 => (256, 144),
        13 => (256, 176),
        14 => (256, 210),
        17 => (256, 74),
        18 => (256, 98),
        21 => (256, 110),
        20 => (32, 18),
        23 => (256, 136),
        _ => (1, 4),
    }
}

// ─── GGUF v3 writer ─────────────────────────────────────────────────────────────────────────

/// A metadata value for [`GgufWriter`]. Only the kinds our loaders and llama.cpp actually
/// consume for text models.
pub enum WVal {
    U32(u32),
    U64(u64),
    F32(f32),
    Bool(bool),
    Str(String),
    StrArray(Vec<String>),
    I32Array(Vec<i32>),
    F32Array(Vec<f32>),
}

/// Minimal standard-conforming GGUF v3 writer: header + KVs + tensor table + aligned data.
/// The output round-trips through [`crate::gguf::GgufFile`] (gated test) and loads in llama.cpp.
pub struct GgufWriter {
    kvs: Vec<(String, WVal)>,
    tensors: Vec<(String, Vec<u64>, u32, Vec<u8>)>, // (name, dims fastest-first, type, bytes)
}

impl Default for GgufWriter {
    fn default() -> Self {
        Self::new()
    }
}

impl GgufWriter {
    pub fn new() -> Self {
        Self { kvs: Vec::new(), tensors: Vec::new() }
    }

    pub fn kv(&mut self, key: &str, v: WVal) {
        self.kvs.push((key.to_string(), v));
    }

    /// Add a tensor. `shape_rows_cols` is the ROW-MAJOR logical shape `[rows, cols]` (or `[n]`);
    /// GGUF stores dims fastest-first, so this reverses internally.
    pub fn tensor(&mut self, name: &str, shape: &[usize], ty: u32, bytes: Vec<u8>) {
        let dims: Vec<u64> = shape.iter().rev().map(|&d| d as u64).collect();
        self.tensors.push((name.to_string(), dims, ty, bytes));
    }

    pub fn write_to(&self, path: &std::path::Path) -> Result<()> {
        let f = std::fs::File::create(path).with_context(|| format!("create {path:?}"))?;
        let mut w = std::io::BufWriter::new(f);
        w.write_all(b"GGUF")?;
        w.write_all(&3u32.to_le_bytes())?;
        w.write_all(&(self.tensors.len() as u64).to_le_bytes())?;
        w.write_all(&(self.kvs.len() as u64).to_le_bytes())?;
        for (k, v) in &self.kvs {
            wstr(&mut w, k)?;
            match v {
                WVal::U32(x) => {
                    w.write_all(&4u32.to_le_bytes())?;
                    w.write_all(&x.to_le_bytes())?;
                }
                WVal::U64(x) => {
                    w.write_all(&10u32.to_le_bytes())?;
                    w.write_all(&x.to_le_bytes())?;
                }
                WVal::F32(x) => {
                    w.write_all(&6u32.to_le_bytes())?;
                    w.write_all(&x.to_le_bytes())?;
                }
                WVal::Bool(x) => {
                    w.write_all(&7u32.to_le_bytes())?;
                    w.write_all(&[u8::from(*x)])?;
                }
                WVal::Str(s) => {
                    w.write_all(&8u32.to_le_bytes())?;
                    wstr(&mut w, s)?;
                }
                WVal::StrArray(a) => {
                    w.write_all(&9u32.to_le_bytes())?;
                    w.write_all(&8u32.to_le_bytes())?;
                    w.write_all(&(a.len() as u64).to_le_bytes())?;
                    for s in a {
                        wstr(&mut w, s)?;
                    }
                }
                WVal::I32Array(a) => {
                    w.write_all(&9u32.to_le_bytes())?;
                    w.write_all(&5u32.to_le_bytes())?;
                    w.write_all(&(a.len() as u64).to_le_bytes())?;
                    for x in a {
                        w.write_all(&x.to_le_bytes())?;
                    }
                }
                WVal::F32Array(a) => {
                    w.write_all(&9u32.to_le_bytes())?;
                    w.write_all(&6u32.to_le_bytes())?;
                    w.write_all(&(a.len() as u64).to_le_bytes())?;
                    for x in a {
                        w.write_all(&x.to_le_bytes())?;
                    }
                }
            }
        }
        // tensor table with aligned offsets
        const ALIGN: u64 = 32;
        let mut off = 0u64;
        let mut offs = Vec::with_capacity(self.tensors.len());
        for (_, _, _, bytes) in &self.tensors {
            offs.push(off);
            off += bytes.len() as u64;
            off = off.div_ceil(ALIGN) * ALIGN;
        }
        for ((name, dims, ty, _), o) in self.tensors.iter().zip(&offs) {
            wstr(&mut w, name)?;
            w.write_all(&(dims.len() as u32).to_le_bytes())?;
            for d in dims {
                w.write_all(&d.to_le_bytes())?;
            }
            w.write_all(&ty.to_le_bytes())?;
            w.write_all(&o.to_le_bytes())?;
        }
        // pad to the data-section alignment, then aligned payloads
        let here = w.stream_position_bytes()?;
        let data_start = here.div_ceil(ALIGN) * ALIGN;
        w.write_all(&vec![0u8; (data_start - here) as usize])?;
        let mut cur = 0u64;
        for ((_, _, _, bytes), o) in self.tensors.iter().zip(&offs) {
            if *o > cur {
                w.write_all(&vec![0u8; (*o - cur) as usize])?;
                cur = *o;
            }
            w.write_all(bytes)?;
            cur += bytes.len() as u64;
        }
        w.flush()?;
        Ok(())
    }
}

fn wstr(w: &mut impl Write, s: &str) -> Result<()> {
    w.write_all(&(s.len() as u64).to_le_bytes())?;
    w.write_all(s.as_bytes())?;
    Ok(())
}

/// `BufWriter` has no stable stream_position without Seek bounds everywhere; track via a tiny
/// trait extension over the underlying file.
trait PosExt {
    fn stream_position_bytes(&mut self) -> Result<u64>;
}
impl PosExt for std::io::BufWriter<std::fs::File> {
    fn stream_position_bytes(&mut self) -> Result<u64> {
        use std::io::Seek;
        self.flush()?;
        Ok(self.get_ref().metadata()?.len().max(self.stream_position()?))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::gguf::dequant_ggml;

    fn ramp(n: usize, seed: u64) -> Vec<f32> {
        let mut s = seed;
        (0..n)
            .map(|_| {
                s ^= s << 13;
                s ^= s >> 7;
                s ^= s << 17;
                ((s >> 40) as u32 as f32 / (1u32 << 24) as f32 - 0.5) * 0.2
            })
            .collect()
    }

    /// Every encoder must round-trip through the ORACLE-VERIFIED decoder with error bounded by
    /// its quantization step — the encode side is checked against the decode side that llama.cpp
    /// itself validated.
    #[test]
    fn encoders_round_trip_through_the_oracle_verified_decoders() {
        let w = ramp(4 * 256, 0x2545F4914F6CDD1D);
        // 17/18 are PARITY-CODED codebooks: the ksign flip may spend up to 2|w| on one element
        // by design (llama.cpp does the same), so their bound only screens gross corruption.
        for (ty, bits) in [(8u32, 8.0f32), (2, 4.0), (6, 5.0), (12, 4.0), (13, 5.0), (14, 6.0), (20, 4.0), (23, 4.0), (18, 1.2), (17, 1.0), (21, 2.8)] {
            let enc = encode_ggml(&w, 256, ty).unwrap();
            let dec = dequant_ggml(&enc, ty, w.len()).unwrap();
            let mut worst = 0f32;
            let amax = w.iter().fold(0f32, |a, &v| a.max(v.abs()));
            for (a, b) in w.iter().zip(&dec) {
                worst = worst.max((a - b).abs());
            }
            // error bound: a few quantization steps of the block range (loose but type-scaled;
            // a LAYOUT error produces errors on the order of amax itself)
            let step = amax / (2f32.powf(bits) - 1.0) * 4.0;
            assert!(
                worst < step.max(1e-4),
                "type {ty}: worst |Δ| {worst} vs step bound {step}"
            );
        }
    }

    /// The K-quant fit must beat a NAIVE range fit on RMSE — the refinement iterations must
    /// actually pay, or the encoder silently ships worse artifacts than llama-quantize.
    #[test]
    fn kquant_fit_beats_naive_range_quantization() {
        let w = ramp(8 * 256, 0x9E3779B97F4A7C15);
        let enc = enc_q4_k(&w, 256, None);
        let dec = dequant_ggml(&enc, 12, w.len()).unwrap();
        let rmse: f64 = w
            .iter()
            .zip(&dec)
            .map(|(a, b)| ((a - b) as f64).powi(2))
            .sum::<f64>()
            .sqrt();
        // naive: per-32-block plain range quantization to 4 bits, no min, no refit
        let mut naive = 0f64;
        for blk in w.chunks_exact(32) {
            let (lo, hi) = blk
                .iter()
                .fold((f32::MAX, f32::MIN), |(l, h), &v| (l.min(v), h.max(v)));
            let s = (hi - lo) / 15.0;
            for &v in blk {
                let q = if s > 0.0 { ((v - lo) / s).round().clamp(0.0, 15.0) } else { 0.0 };
                naive += ((v - (lo + q * s)) as f64).powi(2);
            }
        }
        let naive = naive.sqrt();
        assert!(
            rmse < naive,
            "refined K-quant fit ({rmse:.6}) must beat the naive range fit ({naive:.6})"
        );
    }

    /// imatrix contract, both halves: (a) a uniform imatrix is BYTE-IDENTICAL to no imatrix
    /// (weights scale by a constant, the argmin cannot move); (b) a lopsided imatrix must
    /// cut reconstruction error on the important columns vs the unweighted fit, for every
    /// type that accepts importance.
    #[test]
    fn imatrix_weighting_shifts_error_off_the_important_columns() {
        let cols = 256usize;
        // 4 rows. Importance must vary WITHIN each 32-element block — the fits are invariant
        // to a constant weight factor, so a half-tensor split would (correctly) be a no-op.
        // Even columns: small-magnitude signal, marked important; odd columns: 8x outliers an
        // unweighted fit spends its range on.
        let mut w = Vec::with_capacity(4 * cols);
        let mut s = 0x2545F4914F6CDD1Du64;
        for i in 0..4 * cols {
            s ^= s << 13;
            s ^= s >> 7;
            s ^= s << 17;
            let v = ((s >> 40) as u32 as f32 / (1u32 << 24) as f32 - 0.5) * 0.2;
            w.push(if i % 2 == 1 { v * 8.0 } else { v });
        }
        let uni = vec![1.0f32; cols];
        let lop: Vec<f32> = (0..cols).map(|c| if c % 2 == 0 { 100.0 } else { 1.0 }).collect();
        for ty in [12u32, 13, 14, 2, 6, 20, 23, 18, 17, 21] {
            let plain = encode_ggml(&w, cols, ty).unwrap();
            if matches!(ty, 12 | 13 | 14 | 20 | 23) {
                // K-quant base fits already minimize a weighted SSE; a constant importance
                // factor cannot move the argmin. (Legacy types' base path is the signed-amax
                // CONVENTION, not an argmin — any real search departs from it.)
                let with_uni = encode_ggml_im(&w, cols, ty, Some(&uni)).unwrap();
                assert_eq!(plain, with_uni, "type {ty}: uniform imatrix must be a no-op");
            }
            let with_lop = encode_ggml_im(&w, cols, ty, Some(&lop)).unwrap();
            let d_plain = dequant_ggml(&plain, ty, w.len()).unwrap();
            let d_lop = dequant_ggml(&with_lop, ty, w.len()).unwrap();
            let imp_sse = |d: &[f32]| -> f64 {
                let mut e = 0f64;
                for r in 0..4 {
                    for c in (0..cols).step_by(2) {
                        e += ((w[r * cols + c] - d[r * cols + c]) as f64).powi(2);
                    }
                }
                e
            };
            let (a, b) = (imp_sse(&d_lop), imp_sse(&d_plain));
            assert!(
                a < b,
                "type {ty}: imatrix must cut important-column SSE ({a:.6} !< {b:.6})"
            );
        }
    }

    #[test]
    fn writer_round_trips_through_our_reader() {
        let dir = std::env::temp_dir().join("osfkb_gguf_write_test");
        let _ = std::fs::create_dir_all(&dir);
        let path = dir.join("rt.gguf");
        let mut g = GgufWriter::new();
        g.kv("general.architecture", WVal::Str("llama".into()));
        g.kv("llama.embedding_length", WVal::U32(64));
        g.kv("llama.block_count", WVal::U32(2));
        g.kv("tokenizer.ggml.tokens", WVal::StrArray(vec!["<s>".into(), "a".into()]));
        let w = ramp(4 * 256, 0xD1B54A32D192ED03);
        g.tensor("blk.0.ffn_down.weight", &[4, 256], 12, enc_q4_k(&w, 256, None));
        g.tensor("output_norm.weight", &[64], 0, {
            let v = ramp(64, 7);
            let mut b = Vec::new();
            for x in &v {
                b.extend_from_slice(&x.to_le_bytes());
            }
            b
        });
        g.write_to(&path).unwrap();
        let mut r = crate::gguf::GgufFile::open(&path).unwrap();
        assert_eq!(r.version, 3);
        match r.kvs.get("llama.embedding_length") {
            Some(crate::gguf::GgufValue::U64(64)) => {}
            other => panic!("embedding_length {other:?}"),
        }
        assert_eq!(r.shape("blk.0.ffn_down.weight").unwrap(), vec![4, 256]);
        // decoded tensor values equal a direct decode of the encoder output (byte round-trip)
        let via_file = r.tensor_f32("blk.0.ffn_down.weight").unwrap();
        let direct = dequant_ggml(&enc_q4_k(&w, 256, None), 12, w.len()).unwrap();
        assert_eq!(via_file, direct, "file round-trip must be byte-exact");
        let _ = std::fs::remove_file(&path);
    }
}