// vendor/strata/metal/src/heal_k256_q2_0.cpp - Phase 3 Signed-d AWQ/Lloyd-Max Q2_0 Expert Healer +
// Real-Trace Layerwise Scale Calibration + Monotonic Shared-Expert Ridge Solver for K256_reap_l2_sum
//
// Modes:
//   1. heal_k256_q2_0 --quantize <in-metal.gguf> <out-healed-metal.gguf> [gain_gu=0.50] [gain_d=0.30] [cal_iq3s_pfx]
//      Quantizes all 48 x 256 routed experts to Signed-d AWQ/Lloyd-Max Q2_0 (unlocking the sign bit of fp16 d
//      for a 0-bit polarity switch between {-|d|, 0, +|d|, +2|d|} and {-2|d|, -|d|, 0, +|d|}), weighted by
//      tempered real activation RMS (E[mixed_d^2])^0.5 (beta=0.50) on W_gate/W_up and analytical
//      fourth-root importance (sigma_g * sigma_u)^0.5 ~ (E[h_c^2])^0.25 (beta=0.25) on W_down
//      when cal_iq3s_pfx is provided (falling back to hc_ffn_norm RMS if omitted), and saves <out>.experts.d_bak.bin.
//   2. heal_k256_q2_0 --calibrate-and-heal <in-metal.gguf> <out-healed-metal.gguf> <cal_iq3s_pfx> <yr_q2_bin> <out_writers_bin> [gamma=0.75] [eta_shd=0.15] [ridge_rel=6.0]
//      Calibrates per-layer W_down fp16 d scale alpha_l = cos_l^(1-gamma) / norm_ratio_l directly from the real
//      multi-domain teacher trace (all t_cal tokens x all 256 routed experts, eliminating Gaussian depth overshoot)
//      and solves the ridge regression (full-rank 640/640 on 46/48 layers; L0=430, L2=617) for blk.0..47.ffn_down_shexp.weight with
//      monotonic Q8_0 / IQ4_NL block quantization, writing the 48 modified ffn_down_shexp tensors to <out_writers_bin>.

#ifndef ACCELERATE_NEW_LAPACK
#define ACCELERATE_NEW_LAPACK
#endif
#include "gguf_lite.h"
#include "ref_quant.h"
#include <Accelerate/Accelerate.h>
#include <algorithm>
#include <atomic>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fcntl.h>
#include <string>
#include <sys/stat.h>
#include <thread>
#include <unistd.h>
#include <vector>

static constexpr int NL = 48;
static constexpr int NE = 256;
static constexpr int N  = 2560;
static constexpr int D  = 640;
static constexpr uint64_t ALIGN = 16384;
static constexpr uint64_t Q2_SLICE_BYTES = 460800; // 2560 * 640 / 64 * 18
static constexpr size_t N_BLOCKS = Q2_SLICE_BYTES / 18;

static constexpr int8_t KVALUES_IQ4NL[16] = {
    -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113
};

static void require_io(bool ok, const char* msg) {
    if (!ok) {
        std::perror(msg);
        std::exit(2);
    }
}

static void pread_exact(int fd, void* buf, size_t n, off_t off, const char* msg) {
    uint8_t* p = static_cast<uint8_t*>(buf);
    size_t got = 0;
    while (got < n) {
        ssize_t r = ::pread(fd, p + got, n - got, off + (off_t)got);
        require_io(r > 0, msg);
        got += (size_t)r;
    }
}

static void pwrite_exact(int fd, const void* buf, size_t n, off_t off, const char* msg) {
    const uint8_t* p = static_cast<const uint8_t*>(buf);
    size_t put = 0;
    while (put < n) {
        ssize_t w = ::pwrite(fd, p + put, n - put, off + (off_t)put);
        require_io(w > 0, msg);
        put += (size_t)w;
    }
}

static uint64_t fnv1a_64(const void* data, size_t nbytes) {
    const uint8_t* p = static_cast<const uint8_t*>(data);
    uint64_t h = 14695981039346656037ULL;
    for (size_t i = 0; i < nbytes; ++i) {
        h ^= (uint64_t)p[i];
        h *= 1099511628211ULL;
    }
    return h;
}

static void quantize_q2_0_signed_awq_fast(
    const float* x, size_t n_rows, size_t n_cols, const float* col_imp, uint8_t* out, float gain_exp
) {
    const size_t blocks_per_row = n_cols / 64;
    for (size_t r = 0; r < n_rows; ++r) {
        for (size_t b = 0; b < blocks_per_row; ++b) {
            const float* xb = x + r * n_cols + b * 64;
            const float* wb = col_imp ? (col_imp + b * 64) : nullptr;
            float wxb[64], wi_arr[64];
            float amax = 0.0f, norm_x2 = 0.0f;
            for (int i = 0; i < 64; ++i) {
                const float xi = xb[i];
                const float wi = wb ? wb[i] : 1.0f;
                wi_arr[i] = wi;
                wxb[i] = wi * xi;
                amax = std::max(amax, std::fabs(xi));
                norm_x2 += wxb[i] * xi;
            }
            int best_q[64] = {0};
            uint16_t d16 = 0;
            if (amax > 0.0f && norm_x2 > 0.0f) {
                float best_score = -1.0f, best_qw = 0.0f, best_qq = 1.0f;
                static constexpr float seeds[10] = {
                    0.42f, 0.52f, 0.64f, 0.78f, 0.94f,
                    -0.42f, -0.52f, -0.64f, -0.78f, -0.94f
                };
                for (int s = 0; s < 10; ++s) {
                    float d = amax * seeds[s];
                    for (int iter = 0; iter < 2; ++iter) {
                        if (std::fabs(d) <= 1e-12f) break;
                        const float inv = 1.0f / d;
                        float qw = 0.0f, qq = 0.0f;
                        int q_cur[64];
                        for (int i = 0; i < 64; ++i) {
                            int q = (int)std::lrintf(xb[i] * inv);
                            q = std::max(-1, std::min(2, q));
                            q_cur[i] = q;
                            const float fq = (float)q;
                            qw += fq * wxb[i];
                            qq += wi_arr[i] * (fq * fq);
                        }
                        if (qq <= 0.0f) break;
                        const float d_new = qw / qq;
                        if ( (d > 0.0f ? d_new <= 0.0f : d_new >= 0.0f) ) break;
                        const float score = (qw * qw) / qq;
                        if (score > best_score) {
                            best_score = score;
                            best_qw = qw;
                            best_qq = qq;
                            std::memcpy(best_q, q_cur, sizeof(best_q));
                        }
                        d = d_new;
                    }
                }
                float d_ls = best_qw / best_qq;
                float r2 = best_score / std::max(norm_x2, 1e-20f);
                float gain = (gain_exp == 0.0f) ? 1.0f : std::pow(std::max(r2, 0.25f), -0.5f * gain_exp);
                d16 = ref::f32_to_f16(d_ls * gain);
            }
            uint8_t* dst = out + (r * blocks_per_row + b) * 18;
            std::memcpy(dst, &d16, 2);
            for (int j = 0; j < 16; ++j) {
                uint8_t byte = 0;
                for (int k = 0; k < 4; ++k) {
                    const uint8_t code = (uint8_t)(best_q[j * 4 + k] + 1);
                    byte |= (uint8_t)(code << (2 * k));
                }
                dst[2 + j] = byte;
            }
        }
    }
}

static void quantize_iq4nl_monotonic(const uint8_t* raw_old, const float* w_target, size_t n_elems, uint8_t* raw_new) {
    const size_t n_blks = n_elems / 32;
    std::vector<float> w_old_deq(32);
    for (size_t b = 0; b < n_blks; ++b) {
        const uint8_t* bo = raw_old + b * 18;
        uint8_t* bn = raw_new + b * 18;
        std::memcpy(bn, bo, 18);
        ref::dequant(ref::IQ4_NL, bo, w_old_deq.data(), 32);

        const float* wt = w_target + b * 32;
        float err_old = 0.0f;
        for (int i = 0; i < 32; ++i) {
            float diff = w_old_deq[i] - wt[i];
            err_old += diff * diff;
        }
        uint16_t d_old_u;
        std::memcpy(&d_old_u, bo, 2);
        float d_old = ref::f16_to_f32(d_old_u);
        if (std::fabs(d_old) < 1e-12f) continue;

        float best_err = err_old;
        uint8_t best_blk[18];
        std::memcpy(best_blk, bo, 18);

        // Search 25 fine scale multipliers around d_old so IQ4_NL preserves small ridge updates
        for (int s = -12; s <= 12; ++s) {
            float mult = 1.0f + 0.015f * (float)s;
            uint16_t d_cand_u = ref::f32_to_f16(d_old * mult);
            float d_cand = ref::f16_to_f32(d_cand_u);
            uint8_t idx[32];
            float err_c = 0.0f;
            for (int i = 0; i < 32; ++i) {
                float best_w_err = 1e30f;
                uint8_t best_q = 0;
                for (uint8_t q = 0; q < 16; ++q) {
                    float v = d_cand * (float)KVALUES_IQ4NL[q];
                    float de = (v - wt[i]) * (v - wt[i]);
                    if (de < best_w_err) {
                        best_w_err = de;
                        best_q = q;
                    }
                }
                idx[i] = best_q;
                err_c += best_w_err;
            }
            if (err_c + 1e-11f < best_err) {
                best_err = err_c;
                std::memcpy(best_blk, &d_cand_u, 2);
                for (int j = 0; j < 16; ++j) {
                    best_blk[2 + j] = (uint8_t)((idx[j] & 0x0F) | ((idx[j + 16] & 0x0F) << 4));
                }
            }
        }
        std::memcpy(bn, best_blk, 18);
    }
}

static void quantize_q8_0_monotonic(const uint8_t* raw_old, const float* w_target, size_t n_elems, uint8_t* raw_new) {
    const size_t n_blks = n_elems / 32;
    std::vector<float> w_old_deq(32);
    for (size_t b = 0; b < n_blks; ++b) {
        const uint8_t* bo = raw_old + b * 34;
        uint8_t* bn = raw_new + b * 34;
        std::memcpy(bn, bo, 34);
        ref::dequant(ref::Q8_0, bo, w_old_deq.data(), 32);

        const float* wt = w_target + b * 32;
        float err_old = 0.0f, amax = 0.0f;
        for (int i = 0; i < 32; ++i) {
            float diff = w_old_deq[i] - wt[i];
            err_old += diff * diff;
            amax = std::max(amax, std::fabs(wt[i]));
        }
        uint16_t d_cand_u = ref::f32_to_f16(amax / 127.0f);
        float d_cand = ref::f16_to_f32(d_cand_u);
        if (d_cand <= 0.0f) continue;
        int8_t qs[32];
        float err_c = 0.0f;
        for (int i = 0; i < 32; ++i) {
            int q = (int)std::lrintf(wt[i] / d_cand);
            q = std::max(-127, std::min(127, q));
            qs[i] = (int8_t)q;
            float diff = d_cand * (float)q - wt[i];
            err_c += diff * diff;
        }
        if (err_c + 1e-11f < err_old) {
            std::memcpy(bn, &d_cand_u, 2);
            std::memcpy(bn + 2, qs, 32);
        }
    }
}

static int run_quantize_mode(int argc, char** argv) {
    if (argc < 4) {
        std::fprintf(stderr, "usage: heal_k256_q2_0 --quantize <in-metal.gguf> <out-healed-metal.gguf> [gain_gu=0.50] [gain_d=0.30] [cal_iq3s_pfx]\n");
        return 2;
    }
    const std::string in_gguf = argv[2];
    const std::string out_gguf = argv[3];
    const float gain_gu = (argc > 4) ? std::atof(argv[4]) : 0.50f;
    const float gain_d  = (argc > 5) ? std::atof(argv[5]) : 0.30f;
    const std::string cal_pfx = (argc > 6) ? argv[6] : "";
    const std::string in_exp = in_gguf + ".experts";
    const std::string out_exp = out_gguf + ".experts";
    const std::string bak_d_path = out_exp + ".d_bak.bin";

    gguf::File f;
    std::string err;
    if (!f.open(in_gguf, err)) {
        std::fprintf(stderr, "gguf open failed: %s\n", err.c_str());
        return 1;
    }

    // Check if real teacher activation dump is available for tempered RMS importance
    bool use_real_act = false;
    size_t t_cal = 0;
    int fd_cal_x = -1;
    if (!cal_pfx.empty()) {
        struct stat st{};
        std::string mix_path = cal_pfx + ".mixed";
        if (::stat(mix_path.c_str(), &st) == 0 && st.st_size > 0 && st.st_size % ((size_t)NL * N * sizeof(float)) == 0) {
            t_cal = (size_t)st.st_size / ((size_t)NL * N * sizeof(float));
            fd_cal_x = ::open(mix_path.c_str(), O_RDONLY);
            require_io(fd_cal_x >= 0, "open cal_pfx.mixed");
            use_real_act = (t_cal >= 32);
        }
    }

    std::vector<std::vector<float>> layer_s_in(NL, std::vector<float>(N, 1.0f));
    std::vector<std::vector<float>> layer_s2_raw(NL, std::vector<float>(N, 1.0f));
    if (use_real_act) {
        std::vector<float> x_l(t_cal * N);
        for (int l = 0; l < NL; ++l) {
            pread_exact(fd_cal_x, x_l.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread cal_pfx.mixed");
            for (int d = 0; d < N; ++d) {
                double ms2 = 0.0;
                for (size_t t = 0; t < t_cal; ++t) {
                    double v = x_l[t * N + d];
                    ms2 += v * v;
                }
                double e_x2 = std::max(ms2 / (double)t_cal, 1e-8);
                layer_s2_raw[l][d] = (float)e_x2;
                // Tempered real activation RMS (beta = 0.5: sqrt(E[mixed_d^2]))
                layer_s_in[l][d] = (float)std::sqrt(e_x2);
            }
        }
        ::close(fd_cal_x);
        std::printf("Using tempered real activation RMS (T=%zu, beta=0.50) + per-expert (sigma_g*sigma_u)^0.5 (beta=0.25) for Signed-d Q2_0\n", t_cal);
    } else {
        for (int l = 0; l < NL; ++l) {
            const std::string tname = "blk." + std::to_string(l) + ".hc_ffn_norm.weight";
            const gguf::Tensor* tn = f.find(tname);
            if (tn && tn->type == 0) {
                const float* wn = reinterpret_cast<const float*>(static_cast<const uint8_t*>(f.map) + tn->offset);
                for (int d = 0; d < N; ++d) {
                    float ms2 = 0.0f;
                    for (int c = 0; c < 4; ++c) {
                        const float v = wn[c * N + d];
                        ms2 += v * v;
                    }
                    layer_s_in[l][d] = std::sqrt(std::max(ms2 * 0.25f, 1e-12f));
                }
            }
        }
    }

    const size_t hlen = std::min<size_t>(f.file_bytes, 16 * 1024 * 1024);
    std::vector<uint8_t> hdr(hlen);
    std::memcpy(hdr.data(), f.map, hlen);
    for (int l = 0; l < NL; ++l) {
        for (const char* role : {"gate", "up", "down"}) {
            std::string name = "blk." + std::to_string(l) + ".ffn_" + role + "_exps.weight";
            const uint64_t nlen = name.size();
            std::vector<uint8_t> pat(8 + nlen);
            std::memcpy(pat.data(), &nlen, 8);
            std::memcpy(pat.data() + 8, name.data(), nlen);
            auto it = std::search(hdr.begin(), hdr.end(), pat.begin(), pat.end());
            require_io(it != hdr.end(), ("find tensor in GGUF header: " + name).c_str());
            size_t pos = (size_t)(it - hdr.begin()) + pat.size();
            uint32_t nd = 0;
            std::memcpy(&nd, hdr.data() + pos, 4);
            pos += 4 + (size_t)nd * 8;
            uint32_t new_type = 42;
            std::memcpy(hdr.data() + pos, &new_type, 4);
        }
    }

    int ofd = ::open(out_gguf.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644);
    require_io(ofd >= 0, "open out_gguf");
    fcntl(ofd, F_NOCACHE, 1);
    pwrite_exact(ofd, hdr.data(), hlen, 0, "pwrite out_gguf hdr");
    size_t pos = hlen;
    const size_t chunk = 64ull << 20;
    while (pos < f.file_bytes) {
        const size_t n = std::min(chunk, f.file_bytes - pos);
        pwrite_exact(ofd, (const uint8_t*)f.map + pos, n, (off_t)pos, "pwrite out_gguf body");
        pos += n;
    }
    ::fsync(ofd);
    ::close(ofd);

    int ifd = ::open(in_exp.c_str(), O_RDONLY);
    require_io(ifd >= 0, "open in_exp");
    fcntl(ifd, F_NOCACHE, 1);
    uint64_t in_tab[NL * 4];
    pread_exact(ifd, in_tab, sizeof(in_tab), 64, "pread in_tab");

    const uint64_t exp_data0 = ALIGN;
    uint64_t out_offs[NL];
    uint64_t at = exp_data0;
    const uint64_t out_exp_bytes = 3 * Q2_SLICE_BYTES;
    const uint64_t out_layer_bytes = (uint64_t)NE * out_exp_bytes;
    for (int l = 0; l < NL; ++l) {
        at = (at + ALIGN - 1) & ~(ALIGN - 1);
        out_offs[l] = at;
        at += out_layer_bytes;
    }

    int efd = ::open(out_exp.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644);
    require_io(efd >= 0, "open out_exp");
    fcntl(efd, F_NOCACHE, 1);

    uint8_t strx_hdr[ALIGN];
    std::memset(strx_hdr, 0, sizeof(strx_hdr));
    std::memcpy(strx_hdr, "STRX", 4);
    uint32_t v1 = 1, nl_u = NL, ne_u = NE;
    uint64_t src_sz = f.file_bytes;
    std::memcpy(strx_hdr + 4,  &v1, 4);
    std::memcpy(strx_hdr + 8,  &nl_u, 4);
    std::memcpy(strx_hdr + 12, &ne_u, 4);
    std::memcpy(strx_hdr + 16, &src_sz, 8);
    for (int l = 0; l < NL; ++l) {
        uint64_t row[4] = {out_offs[l], Q2_SLICE_BYTES, Q2_SLICE_BYTES, Q2_SLICE_BYTES};
        std::memcpy(strx_hdr + 64 + l * 32, row, 32);
    }
    pwrite_exact(efd, strx_hdr, ALIGN, 0, "pwrite strx_hdr");

    std::vector<uint16_t> ds_all((size_t)NL * NE * N_BLOCKS);
    const size_t n_elems = (size_t)N * D;
    auto t0 = std::chrono::steady_clock::now();

    for (int l = 0; l < NL; ++l) {
        const std::string b = "blk." + std::to_string(l) + ".";
        int tg = f.find(b + "ffn_gate_exps.weight")->type;
        int tu = f.find(b + "ffn_up_exps.weight")->type;
        int td = f.find(b + "ffn_down_exps.weight")->type;

        const uint64_t in_off = in_tab[l * 4 + 0];
        const uint64_t gs     = in_tab[l * 4 + 1];
        const uint64_t us     = in_tab[l * 4 + 2];
        const uint64_t ds     = in_tab[l * 4 + 3];
        const uint64_t in_exp_step = gs + us + ds;
        const size_t in_layer_bytes = (size_t)NE * in_exp_step;

        std::vector<uint8_t> in_buf(in_layer_bytes);
        pread_exact(ifd, in_buf.data(), in_layer_bytes, (off_t)in_off, "pread in_layer");

        std::vector<uint8_t> out_buf(out_layer_bytes);
        std::atomic<int> next_e{0};
        const int n_threads = 10;
        std::vector<std::thread> threads;
        for (int tid = 0; tid < n_threads; ++tid) {
            threads.emplace_back([&]() {
                std::vector<float> f_gate(n_elems), f_up(n_elems), f_down(n_elems), s_down(D, 1.0f);
                const float* s2_l = layer_s2_raw[l].data();
                for (;;) {
                    int e = next_e.fetch_add(1);
                    if (e >= NE) break;
                    const uint8_t* src = in_buf.data() + (size_t)e * in_exp_step;
                    uint8_t* dst = out_buf.data() + (size_t)e * out_exp_bytes;

                    if (tg == ref::Q2_0) {
                        std::memcpy(dst, src, Q2_SLICE_BYTES);
                        if (use_real_act && td != ref::Q2_0) ref::dequant(tg, src, f_gate.data(), n_elems);
                    } else {
                        ref::dequant(tg, src, f_gate.data(), n_elems);
                        quantize_q2_0_signed_awq_fast(f_gate.data(), D, N, layer_s_in[l].data(), dst, gain_gu);
                    }
                    if (tu == ref::Q2_0) {
                        std::memcpy(dst + Q2_SLICE_BYTES, src + gs, Q2_SLICE_BYTES);
                        if (use_real_act && td != ref::Q2_0) ref::dequant(tu, src + gs, f_up.data(), n_elems);
                    } else {
                        ref::dequant(tu, src + gs, f_up.data(), n_elems);
                        quantize_q2_0_signed_awq_fast(f_up.data(), D, N, layer_s_in[l].data(), dst + Q2_SLICE_BYTES, gain_gu);
                    }
                    uint8_t* qd_dst = dst + 2 * Q2_SLICE_BYTES;
                    if (td == ref::Q2_0) {
                        std::memcpy(qd_dst, src + gs + us, Q2_SLICE_BYTES);
                    } else {
                        ref::dequant(td, src + gs + us, f_down.data(), n_elems);
                        const float* d_imp = nullptr;
                        if (use_real_act) {
                            for (int c = 0; c < D; ++c) {
                                const float* rg = f_gate.data() + (size_t)c * N;
                                const float* ru = f_up.data()   + (size_t)c * N;
                                double vg = 0.0, vu = 0.0;
                                for (int d = 0; d < N; ++d) {
                                    const double s2d = s2_l[d];
                                    vg += (double)rg[d] * rg[d] * s2d;
                                    vu += (double)ru[d] * ru[d] * s2d;
                                }
                                s_down[c] = (float)std::pow(std::max(vg * vu, 1e-16), 0.25);
                            }
                            d_imp = s_down.data();
                        }
                        quantize_q2_0_signed_awq_fast(f_down.data(), N, D, d_imp, qd_dst, gain_d);
                    }
                    uint16_t* d_bak = ds_all.data() + ((size_t)l * NE + e) * N_BLOCKS;
                    for (size_t blk = 0; blk < N_BLOCKS; ++blk) {
                        std::memcpy(&d_bak[blk], qd_dst + blk * 18, 2);
                    }
                }
            });
        }
        for (auto& th : threads) th.join();
        pwrite_exact(efd, out_buf.data(), out_buf.size(), (off_t)out_offs[l], "pwrite out_layer");
        if ((l + 1) % 8 == 0 || l == NL - 1) {
            double el = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();
            std::printf("  Quantized layer %d/%d to Signed-d AWQ Q2_0 (%.1fs elapsed)\n", l + 1, NL, el);
            std::fflush(stdout);
        }
    }

    FILE* fb = std::fopen(bak_d_path.c_str(), "wb");
    require_io(fb != nullptr, "fopen d_bak.bin");
    require_io(std::fwrite(ds_all.data(), sizeof(uint16_t), ds_all.size(), fb) == ds_all.size(), "fwrite d_bak.bin");
    std::fclose(fb);

    // Save .quant.meta fingerprint over (gain_gu, gain_d, cal_pfx, t_cal)
    const std::string qmeta_path = out_exp + ".quant.meta";
    FILE* fqm = std::fopen(qmeta_path.c_str(), "wb");
    require_io(fqm != nullptr, ("fopen " + qmeta_path).c_str());
    uint32_t gu_bits = 0, gd_bits = 0;
    std::memcpy(&gu_bits, &gain_gu, 4);
    std::memcpy(&gd_bits, &gain_d, 4);
    const uint64_t pfx_hash = fnv1a_64(cal_pfx.data(), cal_pfx.size());
    const uint64_t qmeta[3] = {((uint64_t)gu_bits << 32) | gd_bits, pfx_hash, (uint64_t)t_cal};
    require_io(std::fwrite(qmeta, sizeof(uint64_t), 3, fqm) == 3, "fwrite quant.meta");
    std::fclose(fqm);

    ::fsync(efd);
    ::close(efd);
    ::close(ifd);
    return 0;
}

static int run_calibrate_and_heal_mode(int argc, char** argv) {
    if (argc < 7) {
        std::fprintf(stderr, "usage: heal_k256_q2_0 --calibrate-and-heal <in-metal.gguf> <out-healed-metal.gguf> <cal_iq3s_pfx> <yr_q2_bin> <out_writers_bin> [gamma=0.75] [eta_shd=0.15] [ridge_rel=6.0]\n");
        return 2;
    }
    const std::string orig_gguf = argv[2];
    const std::string healed_gguf = argv[3];
    const std::string cal_pfx = argv[4];
    const std::string yr_q2_bin = argv[5];
    const std::string out_writers_bin = argv[6];
    const float gamma     = (argc > 7) ? std::atof(argv[7]) : 0.75f;
    const float eta_shd   = (argc > 8) ? std::atof(argv[8]) : 0.15f;
    const float ridge_rel = (argc > 9) ? std::atof(argv[9]) : 6.0f;

    const std::string healed_exp = healed_gguf + ".experts";
    const std::string bak_d_path = healed_exp + ".d_bak.bin";

    gguf::File gf;
    std::string err;
    if (!gf.open(orig_gguf, err)) {
        std::fprintf(stderr, "gguf open failed: %s\n", err.c_str());
        return 1;
    }

    int fd_heal = ::open(healed_exp.c_str(), O_RDWR);
    require_io(fd_heal >= 0, "open healed_exp");
    uint64_t tab_heal[NL * 4];
    pread_exact(fd_heal, tab_heal, sizeof(tab_heal), 64, "pread tab_heal");

    FILE* fb = std::fopen(bak_d_path.c_str(), "rb");
    require_io(fb != nullptr, ("fopen " + bak_d_path).c_str());
    std::vector<uint16_t> ds_all((size_t)NL * NE * N_BLOCKS);
    require_io(std::fread(ds_all.data(), sizeof(uint16_t), ds_all.size(), fb) == ds_all.size(), "fread d_bak.bin");
    std::fclose(fb);

    // Infer t_cal dynamically from (cal_pfx + ".mixed") file size and cross-check .yr, .ids, .wts, .hh sizes
    struct stat st_mix{}, st_yr{}, st_ids{}, st_wts{}, st_hh{};
    require_io(::stat((cal_pfx + ".mixed").c_str(), &st_mix) == 0 && st_mix.st_size > 0, "stat cal_pfx.mixed");
    require_io(st_mix.st_size % ((size_t)NL * N * sizeof(float)) == 0, "cal_pfx.mixed size divisible by NL*N*4");
    const size_t t_cal = (size_t)st_mix.st_size / ((size_t)NL * N * sizeof(float));
    require_io(::stat((cal_pfx + ".yr").c_str(), &st_yr) == 0 && (size_t)st_yr.st_size == (size_t)NL * t_cal * N * sizeof(float), "cal_pfx.yr size match");
    require_io(::stat((cal_pfx + ".ids").c_str(), &st_ids) == 0 && (size_t)st_ids.st_size == (size_t)NL * t_cal * 10 * sizeof(int32_t), "cal_pfx.ids size match");
    require_io(::stat((cal_pfx + ".wts").c_str(), &st_wts) == 0 && (size_t)st_wts.st_size == (size_t)NL * t_cal * 10 * sizeof(float), "cal_pfx.wts size match");
    require_io(::stat((cal_pfx + ".hh").c_str(), &st_hh) == 0 && (size_t)st_hh.st_size == (size_t)NL * t_cal * D * sizeof(uint16_t), "cal_pfx.hh size match");

    int fd_mix    = ::open((cal_pfx + ".mixed").c_str(), O_RDONLY);
    int fd_yr_ref = ::open((cal_pfx + ".yr").c_str(),    O_RDONLY);
    require_io(fd_mix >= 0 && fd_yr_ref >= 0, "open cal_pfx .mixed/.yr");

    // Compute cache fingerprint over d_bak.bin + L0 & L47 of .mixed + L0/L24/L47 gate/up slices of healed_exp + t_cal
    std::vector<float> mix_l0_l47(2 * t_cal * N);
    pread_exact(fd_mix, mix_l0_l47.data(), t_cal * N * sizeof(float), 0, "pread mix_l0");
    pread_exact(fd_mix, mix_l0_l47.data() + t_cal * N, t_cal * N * sizeof(float), (off_t)((size_t)(NL - 1) * t_cal * N * sizeof(float)), "pread mix_l47");
    std::vector<uint8_t> gu_sample(3 * 2 * Q2_SLICE_BYTES);
    pread_exact(fd_heal, gu_sample.data() + 0 * 2 * Q2_SLICE_BYTES, 2 * Q2_SLICE_BYTES, (off_t)tab_heal[0 * 4 + 0],  "pread gu_l0");
    pread_exact(fd_heal, gu_sample.data() + 1 * 2 * Q2_SLICE_BYTES, 2 * Q2_SLICE_BYTES, (off_t)tab_heal[24 * 4 + 0], "pread gu_l24");
    pread_exact(fd_heal, gu_sample.data() + 2 * 2 * Q2_SLICE_BYTES, 2 * Q2_SLICE_BYTES, (off_t)tab_heal[47 * 4 + 0], "pread gu_l47");
    const uint64_t dbak_hash = fnv1a_64(ds_all.data(), ds_all.size() * sizeof(uint16_t));
    const uint64_t mix_hash  = fnv1a_64(mix_l0_l47.data(), mix_l0_l47.size() * sizeof(float));
    const uint64_t gu_hash   = fnv1a_64(gu_sample.data(), gu_sample.size());
    const uint64_t want_meta[4] = {dbak_hash, mix_hash, gu_hash, (uint64_t)t_cal};
    const std::string meta_path = yr_q2_bin + ".meta";

    bool cache_valid = false;
    struct stat st_q2{};
    if (::stat(yr_q2_bin.c_str(), &st_q2) == 0 && (size_t)st_q2.st_size == (size_t)NL * t_cal * N * sizeof(float)) {
        FILE* fm = std::fopen(meta_path.c_str(), "rb");
        if (fm) {
            uint64_t got_meta[4] = {0, 0, 0, 0};
            if (std::fread(got_meta, sizeof(uint64_t), 4, fm) == 4 &&
                std::memcmp(got_meta, want_meta, sizeof(want_meta)) == 0) {
                cache_valid = true;
            }
            std::fclose(fm);
        }
    }

    if (!cache_valid) {
        int fd_ids = ::open((cal_pfx + ".ids").c_str(), O_RDONLY);
        int fd_wts = ::open((cal_pfx + ".wts").c_str(), O_RDONLY);
        int fd_q2w = ::open(yr_q2_bin.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644);
        require_io(fd_ids >= 0 && fd_wts >= 0 && fd_q2w >= 0, "open cal_pfx .ids/.wts or yr_q2_bin");

        std::atomic<int> nl_q2{0};
        std::vector<std::thread> q2_ths;
        for (int tid = 0; tid < 8; ++tid) {
            q2_ths.emplace_back([&]() {
                const size_t n_elems = (size_t)N * D;
                std::vector<float> x(t_cal * N), yr_q2(t_cal * N, 0.0f);
                std::vector<int32_t> ids(t_cal * 10);
                std::vector<float> wts(t_cal * 10);
                std::vector<uint8_t> in_h((size_t)NE * 3 * Q2_SLICE_BYTES);
                std::vector<float> wg(n_elems), wu(n_elems), wd(n_elems);
                for (;;) {
                    int l = nl_q2.fetch_add(1);
                    if (l >= NL) break;
                    pread_exact(fd_mix,  x.data(),    t_cal * N * sizeof(float),    (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread mix");
                    pread_exact(fd_ids,  ids.data(),  t_cal * 10 * sizeof(int32_t), (off_t)((size_t)l * t_cal * 10 * sizeof(int32_t)), "pread ids");
                    pread_exact(fd_wts,  wts.data(),  t_cal * 10 * sizeof(float),   (off_t)((size_t)l * t_cal * 10 * sizeof(float)), "pread wts");
                    pread_exact(fd_heal, in_h.data(), in_h.size(),                  (off_t)tab_heal[l * 4 + 0], "pread in_h");
                    for (int e = 0; e < NE; ++e) {
                        uint8_t* qd_s = in_h.data() + ((size_t)e * 3 + 2) * Q2_SLICE_BYTES;
                        const uint16_t* d_b = ds_all.data() + ((size_t)l * NE + e) * N_BLOCKS;
                        for (size_t blk = 0; blk < N_BLOCKS; ++blk) std::memcpy(qd_s + blk * 18, &d_b[blk], 2);
                    }
                    std::fill(yr_q2.begin(), yr_q2.end(), 0.0f);
                    std::vector<std::vector<std::pair<int, float>>> exp_toks(NE);
                    for (size_t t = 0; t < t_cal; ++t) {
                        for (int k = 0; k < 10; ++k) {
                            int e = ids[t * 10 + k];
                            float w = wts[t * 10 + k];
                            if (e >= 0 && e < NE && w > 0.0f) exp_toks[e].push_back({(int)t, w});
                        }
                    }
                    std::vector<float> x_e, g_e, u_e, y_e;
                    for (int e = 0; e < NE; ++e) {
                        const auto& hits = exp_toks[e];
                        if (hits.empty()) continue;
                        int m = (int)hits.size();
                        x_e.resize((size_t)m * N); g_e.resize((size_t)m * D); u_e.resize((size_t)m * D); y_e.resize((size_t)m * N);
                        for (int i = 0; i < m; ++i) std::memcpy(x_e.data() + (size_t)i * N, x.data() + (size_t)hits[i].first * N, N * sizeof(float));
                        const uint8_t* rh = in_h.data() + (size_t)e * 3 * Q2_SLICE_BYTES;
                        ref::dequant(ref::Q2_0, rh, wg.data(), n_elems);
                        ref::dequant(ref::Q2_0, rh + Q2_SLICE_BYTES, wu.data(), n_elems);
                        ref::dequant(ref::Q2_0, rh + 2 * Q2_SLICE_BYTES, wd.data(), n_elems);
                        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m, D, N, 1.0f, x_e.data(), N, wg.data(), N, 0.0f, g_e.data(), D);
                        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m, D, N, 1.0f, x_e.data(), N, wu.data(), N, 0.0f, u_e.data(), D);
                        for (int i = 0; i < m * D; ++i) { float gv = g_e[i]; g_e[i] = (gv / (1.0f + std::exp(-gv))) * u_e[i]; }
                        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m, N, D, 1.0f, g_e.data(), D, wd.data(), D, 0.0f, y_e.data(), N);
                        for (int i = 0; i < m; ++i) {
                            float* dst = yr_q2.data() + (size_t)hits[i].first * N;
                            const float* src = y_e.data() + (size_t)i * N;
                            float w = hits[i].second;
                            for (int c = 0; c < N; ++c) dst[c] += w * src[c];
                        }
                    }
                    pwrite_exact(fd_q2w, yr_q2.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pwrite yr_q2");
                }
            });
        }
        for (auto& th : q2_ths) th.join();
        ::close(fd_q2w); ::close(fd_wts); ::close(fd_ids);

        FILE* fm = std::fopen(meta_path.c_str(), "wb");
        require_io(fm != nullptr, "fopen yr_q2_bin.meta");
        require_io(std::fwrite(want_meta, sizeof(uint64_t), 4, fm) == 4, "fwrite yr_q2_bin.meta");
        std::fclose(fm);
    }

    int fd_yr_q2 = ::open(yr_q2_bin.c_str(), O_RDONLY);
    require_io(fd_yr_q2 >= 0, "open yr_q2_bin");

    // Part A (Real-Trace Layerwise Scale Calibration) + Part B (Full-Rank Shared-Expert Ridge Regression)
    FILE* fout = std::fopen(out_writers_bin.c_str(), "wb");
    require_io(fout != nullptr, "fopen out_writers_bin");
    uint32_t n_tensors = (eta_shd > 0.0f) ? NL : 0;
    require_io(std::fwrite(&n_tensors, sizeof(uint32_t), 1, fout) == 1, "fwrite n_tensors");

    auto write_record = [&](const std::string& name, uint32_t gtype, const void* data, uint32_t nbytes) {
        uint32_t nlen = (uint32_t)name.size();
        require_io(std::fwrite(&nlen, 4, 1, fout) == 1, "fwrite nlen");
        require_io(std::fwrite(name.data(), 1, nlen, fout) == nlen, "fwrite name");
        require_io(std::fwrite(&gtype, 4, 1, fout) == 1, "fwrite gtype");
        require_io(std::fwrite(&nbytes, 4, 1, fout) == 1, "fwrite nbytes");
        require_io(std::fwrite(data, 1, nbytes, fout) == nbytes, "fwrite data");
    };

    std::vector<float> layer_alpha(NL, 1.0f);
    std::vector<uint8_t> slice(Q2_SLICE_BYTES);
    double tot_ref2 = 0.0, tot_dy2_unscaled = 0.0, tot_dy2_before = 0.0, tot_dy2_after = 0.0;

    for (int l = 0; l < NL; ++l) {
        std::vector<float> x(t_cal * N), yr_r(t_cal * N), yr_q(t_cal * N), dy(t_cal * N);
        pread_exact(fd_mix,    x.data(),    t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread mix");
        pread_exact(fd_yr_ref, yr_r.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread yr_ref");
        pread_exact(fd_yr_q2,  yr_q.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread yr_q2");

        double sum_ref2 = 0.0, sum_q2 = 0.0, sum_dot = 0.0;
        for (size_t i = 0; i < t_cal * N; ++i) {
            double r = yr_r[i], q = yr_q[i];
            sum_ref2 += r * r;
            sum_q2   += q * q;
            sum_dot  += r * q;
            double du = r - q;
            tot_dy2_unscaled += du * du;
        }
        tot_ref2 += sum_ref2;

        float alpha_l = 1.0f;
        if (gamma >= 0.0f && sum_ref2 > 0.0 && sum_q2 > 0.0) {
            double norm_ratio = std::sqrt(sum_q2 / sum_ref2);
            double cos_sim = std::max(sum_dot / (std::sqrt(sum_ref2) * std::sqrt(sum_q2)), 0.1);
            alpha_l = (float)(std::pow(cos_sim, 1.0 - (double)gamma) / norm_ratio);
        }
        layer_alpha[l] = alpha_l;

        // Apply alpha_l to W_down fp16 d headers in healed_exp
        uint64_t off_h = tab_heal[l * 4 + 0];
        for (int e = 0; e < NE; ++e) {
            const uint16_t* d_bak = ds_all.data() + ((size_t)l * NE + e) * N_BLOCKS;
            uint64_t d_off = off_h + (uint64_t)e * (3 * Q2_SLICE_BYTES) + 2 * Q2_SLICE_BYTES;
            pread_exact(fd_heal, slice.data(), Q2_SLICE_BYTES, (off_t)d_off, "pread slice");
            for (size_t b = 0; b < N_BLOCKS; ++b) {
                float d = ref::f16_to_f32(d_bak[b]);
                uint16_t d_new = ref::f32_to_f16(d * alpha_l);
                std::memcpy(slice.data() + b * 18, &d_new, 2);
            }
            pwrite_exact(fd_heal, slice.data(), Q2_SLICE_BYTES, (off_t)d_off, "pwrite slice");
        }

        for (size_t i = 0; i < t_cal * N; ++i) {
            dy[i] = yr_r[i] - alpha_l * yr_q[i];
            tot_dy2_before += (double)dy[i] * dy[i];
        }

        if (eta_shd <= 0.0f) {
            tot_dy2_after = tot_dy2_before;
            continue;
        }

        const std::string b = "blk." + std::to_string(l) + ".";
        const auto* tn_shg = gf.find(b + "ffn_gate_shexp.weight");
        const auto* tn_shu = gf.find(b + "ffn_up_shexp.weight");
        const auto* tn_shd = gf.find(b + "ffn_down_shexp.weight");
        const auto* tn_sg  = gf.find(b + "ffn_gate_inp_shexp.weight");

        std::vector<float> w_shg(D * N), w_shu(D * N), w_shd(N * D), w_sg(N);
        ref::dequant(tn_shg->type, (const uint8_t*)gf.map + tn_shg->offset, w_shg.data(), D * N);
        ref::dequant(tn_shu->type, (const uint8_t*)gf.map + tn_shu->offset, w_shu.data(), D * N);
        ref::dequant(tn_shd->type, (const uint8_t*)gf.map + tn_shd->offset, w_shd.data(), N * D);
        std::memcpy(w_sg.data(), (const uint8_t*)gf.map + tn_sg->offset, N * sizeof(float));

        std::vector<float> g_sh(t_cal * D), u_sh(t_cal * D), z_sh(t_cal * D), sg_l(t_cal);
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (int)t_cal, D, N, 1.0f, x.data(), N, w_shg.data(), N, 0.0f, g_sh.data(), D);
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (int)t_cal, D, N, 1.0f, x.data(), N, w_shu.data(), N, 0.0f, u_sh.data(), D);
        cblas_sgemv(CblasRowMajor, CblasNoTrans, (int)t_cal, N, 1.0f, x.data(), N, w_sg.data(), 1, 0.0f, sg_l.data(), 1);

        for (size_t t = 0; t < t_cal; ++t) {
            float sig = 1.0f / (1.0f + std::exp(-sg_l[t]));
            for (int c = 0; c < D; ++c) {
                float gv = g_sh[t * D + c];
                z_sh[t * D + c] = sig * (gv / (1.0f + std::exp(-gv))) * u_sh[t * D + c];
            }
        }

        std::vector<float> G(D * D, 0.0f), B(D * N, 0.0f);
        cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, D, D, (int)t_cal, 1.0f, z_sh.data(), D, z_sh.data(), D, 0.0f, G.data(), D);
        cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, D, N, (int)t_cal, 1.0f, z_sh.data(), D, dy.data(), N, 0.0f, B.data(), N);

        double trace_G = 0.0;
        for (int c = 0; c < D; ++c) trace_G += G[c * D + c];
        float lam = (float)(ridge_rel * (trace_G / D));
        for (int c = 0; c < D; ++c) G[c * D + c] += lam;

        std::vector<double> Gd(D * D), L_chol(D * D, 0.0);
        for (int i = 0; i < D * D; ++i) Gd[i] = G[i];
        for (int i = 0; i < D; ++i) {
            for (int j = 0; j <= i; ++j) {
                double s = Gd[i * D + j];
                for (int k = 0; k < j; ++k) s -= L_chol[i * D + k] * L_chol[j * D + k];
                if (i == j) L_chol[i * D + j] = std::sqrt(std::max(s, 1e-12));
                else L_chol[i * D + j] = s / L_chol[j * D + j];
            }
        }
        std::vector<float> w_shd_target = w_shd;
        for (int r = 0; r < N; ++r) {
            double y_fwd[D], x_bwd[D];
            for (int i = 0; i < D; ++i) {
                double s = B[i * N + r];
                for (int k = 0; k < i; ++k) s -= L_chol[i * D + k] * y_fwd[k];
                y_fwd[i] = s / L_chol[i * D + i];
            }
            for (int i = D - 1; i >= 0; --i) {
                double s = y_fwd[i];
                for (int k = i + 1; k < D; ++k) s -= L_chol[k * D + i] * x_bwd[k];
                x_bwd[i] = s / L_chol[i * D + i];
            }
            for (int c = 0; c < D; ++c) {
                w_shd_target[r * D + c] += eta_shd * (float)x_bwd[c];
            }
        }

        const uint8_t* raw_shd_old = (const uint8_t*)gf.map + tn_shd->offset;
        std::vector<uint8_t> raw_shd_new(tn_shd->bytes);
        if (tn_shd->type == ref::Q8_0) {
            quantize_q8_0_monotonic(raw_shd_old, w_shd_target.data(), N * D, raw_shd_new.data());
        } else {
            quantize_iq4nl_monotonic(raw_shd_old, w_shd_target.data(), N * D, raw_shd_new.data());
        }

        std::vector<float> w_shd_new_deq(N * D), delta_w(N * D), pred_dy(t_cal * N);
        ref::dequant(tn_shd->type, raw_shd_new.data(), w_shd_new_deq.data(), N * D);
        for (int i = 0; i < N * D; ++i) delta_w[i] = w_shd_new_deq[i] - w_shd[i];
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (int)t_cal, N, D, 1.0f, z_sh.data(), D, delta_w.data(), D, 0.0f, pred_dy.data(), N);
        for (size_t i = 0; i < t_cal * N; ++i) {
            double rem = (double)dy[i] - (double)pred_dy[i];
            tot_dy2_after += rem * rem;
        }

        write_record(b + "ffn_down_shexp.weight", (uint32_t)tn_shd->type, raw_shd_new.data(), (uint32_t)raw_shd_new.size());
    }

    std::fclose(fout);
    ::fsync(fd_heal);
    ::close(fd_heal);
    ::close(fd_mix);
    ::close(fd_yr_ref);
    ::close(fd_yr_q2);

    std::printf("Calibrated 48 layers on real trace (T=%zu, gamma=%.2f, L0=%.3f L24=%.3f L47=%.3f) & solved %u ffn_down_shexp tensors (MoE MSE %.2f -> %.2f -> %.2f, %.2f%%)\n",
                t_cal, gamma, layer_alpha[0], layer_alpha[24], layer_alpha[47], n_tensors,
                tot_dy2_unscaled, tot_dy2_before, tot_dy2_after, (tot_dy2_after / tot_dy2_unscaled - 1.0) * 100.0);
    return 0;
}

int main(int argc, char** argv) {
    if (argc >= 2 && std::string(argv[1]) == "--quantize") {
        return run_quantize_mode(argc, argv);
    }
    if (argc >= 2 && std::string(argv[1]) == "--calibrate-and-heal") {
        return run_calibrate_and_heal_mode(argc, argv);
    }
    if (argc >= 3) {
        std::vector<char*> args = {argv[0], (char*)"--quantize"};
        for (int i = 1; i < argc; ++i) args.push_back(argv[i]);
        return run_quantize_mode((int)args.size(), args.data());
    }
    std::fprintf(stderr, "usage: heal_k256_q2_0 [--quantize | --calibrate-and-heal] ...\n");
    return 2;
}