#define ACCELERATE_NEW_LAPACK
#include "vendor/strata/metal/src/gguf_lite.h"
#include "vendor/strata/metal/src/ref_quant.h"
#include <Accelerate/Accelerate.h>
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <map>
#include <set>
#include <string>
#include <vector>

static constexpr int N = 2560, HC = 4, LR = 320, F = 640, V = 248320;

static std::vector<char> rdall(const std::string& p) {
    FILE* f = std::fopen(p.c_str(), "rb");
    if (!f) { std::perror(p.c_str()); std::exit(2); }
    std::fseek(f, 0, SEEK_END);
    long n = std::ftell(f);
    std::fseek(f, 0, SEEK_SET);
    std::vector<char> b(n);
    if (std::fread(b.data(), 1, n, f) != (size_t)n) { std::perror("fread"); std::exit(2); }
    std::fclose(f);
    return b;
}

static std::vector<int> rdtoks(const std::string& p) {
    std::vector<int> v;
    FILE* f = std::fopen(p.c_str(), "r");
    if (!f) { std::perror(p.c_str()); std::exit(2); }
    int x;
    while (std::fscanf(f, "%d", &x) == 1) v.push_back(x);
    std::fclose(f);
    return v;
}

struct TraceData {
    int T = 0;
    std::vector<int> toks;
    std::vector<float> R;     // [T, HC, N]
    std::vector<float> inj;   // [T, HC]
    std::vector<float> z;     // [T, F] = sigmoid(sgv) * hh
    std::vector<float> mix0;  // [T, N] at dW = 0
    std::vector<float> gam;   // [T, N] diagonal Jacobian d(mix)/d(d) at dW = 0
};

static void compute_mix_batch(
    int T,
    const float* R,
    const float* inj,
    const float* d_in, // [T, N] or nullptr
    const std::vector<float>& wn,
    const std::vector<float>& wdn,
    const std::vector<float>& wup,
    float* mix_out,
    float* gam_out = nullptr
) {
    std::vector<float> xn((size_t)T * HC * N);
    std::vector<float> rs_all((size_t)T * HC);
    std::vector<float> gg_all((size_t)T * HC);
    for (int t = 0; t < T; ++t) {
        const float* dt = d_in ? &d_in[(size_t)t * N] : nullptr;
        for (int c = 0; c < HC; ++c) {
            float gg = 2.0f / (1.0f + std::exp(-inj[t * HC + c] / 4.0f));
            gg_all[t * HC + c] = gg;
            const float* rc = &R[((size_t)t * HC + c) * N];
            float* xnc = &xn[((size_t)t * HC + c) * N];
            double ss = 0.0;
            for (int j = 0; j < N; ++j) {
                float v = rc[j] + (dt ? gg * dt[j] : 0.0f);
                xnc[j] = v;
                ss += (double)v * v;
            }
            float rs = 1.0f / std::sqrt((float)(ss / N) + 1e-6f);
            rs_all[t * HC + c] = rs;
            const float* wnc = &wn[c * N];
            for (int j = 0; j < N; ++j) xnc[j] *= wnc[j] * rs;
        }
    }
    std::vector<float> lo((size_t)T * LR);
    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, T, LR, HC * N, 0.25f, xn.data(), HC * N, wdn.data(), HC * N, 0.0f, lo.data(), LR);
    for (size_t i = 0; i < lo.size(); ++i) {
        float z = lo[i];
        lo[i] = z / (1.0f + std::exp(-z));
    }
    std::vector<float> gup((size_t)T * HC * N);
    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, T, HC * N, LR, 1.0f, lo.data(), LR, wup.data(), LR, 0.0f, gup.data(), HC * N);
    for (int t = 0; t < T; ++t) {
        float* mix = &mix_out[(size_t)t * N];
        float* gm  = gam_out ? &gam_out[(size_t)t * N] : nullptr;
        for (int j = 0; j < N; ++j) {
            float m = 0.0f, gs = 0.0f;
            for (int c = 0; c < HC; ++c) {
                float s = gup[((size_t)t * HC + c) * N + j];
                float gate = 1.0f / (1.0f + std::exp(-s));
                m += xn[((size_t)t * HC + c) * N + j] * gate;
                if (gm) gs += gg_all[t * HC + c] * rs_all[t * HC + c] * wn[c * N + j] * gate;
            }
            mix[j] = m * 0.25f;
            if (gm) gm[j] = gs * 0.25f;
        }
    }
}

static void load_trace(
    const std::string& pfx,
    const std::string& tok_file,
    const std::vector<float>& wn,
    const std::vector<float>& wdn,
    const std::vector<float>& wup,
    TraceData& tr
) {
    tr.toks = rdtoks(tok_file);
    auto Rb = rdall(pfx + ".R_post47");
    auto Ib = rdall(pfx + ".inj47");
    auto Hb = rdall(pfx + ".hh47");
    auto Sb = rdall(pfx + ".sgv47");
    tr.T = (int)(Rb.size() / (HC * N * sizeof(float)));
    tr.R.resize((size_t)tr.T * HC * N);
    tr.inj.resize((size_t)tr.T * HC);
    std::memcpy(tr.R.data(), Rb.data(), Rb.size());
    std::memcpy(tr.inj.data(), Ib.data(), Ib.size());

    const __fp16* hh = (const __fp16*)Hb.data();
    const float* sg = (const float*)Sb.data();
    tr.z.resize((size_t)tr.T * F);
    for (int t = 0; t < tr.T; ++t) {
        float sgs = 1.0f / (1.0f + std::exp(-sg[t]));
        for (int k = 0; k < F; ++k) tr.z[(size_t)t * F + k] = sgs * (float)hh[(size_t)t * F + k];
    }
    tr.mix0.resize((size_t)tr.T * N);
    tr.gam.resize((size_t)tr.T * N);
    compute_mix_batch(tr.T, tr.R.data(), tr.inj.data(), nullptr, wn, wdn, wup, tr.mix0.data(), tr.gam.data());
}


struct FtpoRow {
    int split, phrase_id, t_pos, rej_id;
    std::vector<int> alt_ids;
};
struct ColRow {
    int split, ban_idx, t_pos, target_id;
};
struct SepHoRow {
    int phrase_id, t_pos, rej_id;
};

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;
        for (int i = 0; i < 32; ++i) {
            float diff = w_old_deq[i] - wt[i];
            err_old += diff * diff;
        }
        uint8_t cand_blk[34];
        ref::quantize_q8_0(wt, 32, cand_blk);
        std::vector<float> w_cand_deq(32);
        ref::dequant(ref::Q8_0, cand_blk, w_cand_deq.data(), 32);
        float err_c = 0.0f;
        for (int i = 0; i < 32; ++i) {
            float diff = w_cand_deq[i] - wt[i];
            err_c += diff * diff;
        }
        if (err_c + 1e-11f < err_old) std::memcpy(bn, cand_blk, 34);
    }
}

int main() {
    gguf::File g;
    std::string err;
    if (!g.open("models/K256_reap_l2_sum/K256-q2_0-healed-metal.gguf", err)) return 1;
    auto deq = [&](const char* n, std::vector<float>& out) {
        auto* t = g.find(n);
        size_t ne = 1;
        for (auto d : t->shape) ne *= d;
        out.resize(ne);
        ref::dequant(t->type, (const uint8_t*)g.map + t->offset, out.data(), ne);
    };
    auto q8rt = [&](std::vector<float>& w) {
        std::vector<uint8_t> q(w.size() / 32 * 34);
        ref::quantize_q8_0(w.data(), w.size(), q.data());
        ref::dequant(ref::Q8_0, q.data(), w.data(), w.size());
    };
    std::vector<float> wn, wdn, wup, wb;
    deq("output_hc_norm.weight", wn);
    deq("output_hc_down.weight", wdn); q8rt(wdn);
    deq("output_hc_up.weight", wup);   q8rt(wup);
    deq("blk.47.ffn_down_shexp.weight", wb);
    auto* t_sd47 = g.find("blk.47.ffn_down_shexp.weight");
    const uint8_t* raw_sd47_orig = (const uint8_t*)g.map + t_sd47->offset;

    TraceData tr_cal, tr_p4, tr_col, tr_ho;
    load_trace("models/K256_reap_l2_sum/cal_cache/cal_p3_1024", "data/m5_val_prompts/cal_disjoint_1024tok.txt", wn, wdn, wup, tr_cal);
    load_trace("models/K256_reap_l2_sum/cal_cache/p4_ftpo_sft_1024", "data/m5_val_prompts/phase4_ftpo_sft_1024tok.txt", wn, wdn, wup, tr_p4);
    load_trace("models/K256_reap_l2_sum/cal_cache/p4_collateral_1640", "data/m5_val_prompts/phase4_collateral_1024tok.txt", wn, wdn, wup, tr_col);
    load_trace("models/K256_reap_l2_sum/cal_cache/p4_ftpo_heldout_1024", "data/m5_val_prompts/phase4_ftpo_heldout_1024tok.txt", wn, wdn, wup, tr_ho);

    FILE* fm = std::fopen("/tmp/p4_meta_table.txt", "r");
    int n_ftpo = 0; std::fscanf(fm, "%d", &n_ftpo);
    std::vector<FtpoRow> ftpo(n_ftpo);
    std::set<int> all_p4_slop_set, train_p4_slop_set;
    for (int i = 0; i < n_ftpo; ++i) {
        int na = 0;
        std::fscanf(fm, "%d %d %d %d %d", &ftpo[i].split, &ftpo[i].phrase_id, &ftpo[i].t_pos, &ftpo[i].rej_id, &na);
        ftpo[i].alt_ids.resize(na);
        for (int k = 0; k < na; ++k) std::fscanf(fm, "%d", &ftpo[i].alt_ids[k]);
        all_p4_slop_set.insert(ftpo[i].t_pos);
        if (ftpo[i].split == 0) train_p4_slop_set.insert(ftpo[i].t_pos);
    }
    int n_col = 0; std::fscanf(fm, "%d", &n_col);
    std::vector<ColRow> col(n_col);
    for (int i = 0; i < n_col; ++i) std::fscanf(fm, "%d %d %d %d", &col[i].split, &col[i].ban_idx, &col[i].t_pos, &col[i].target_id);
    int n_sep = 0; std::fscanf(fm, "%d", &n_sep);
    std::vector<SepHoRow> sep_ho(n_sep);
    std::set<int> sep_ho_set;
    for (int i = 0; i < n_sep; ++i) {
        std::fscanf(fm, "%d %d %d", &sep_ho[i].phrase_id, &sep_ho[i].t_pos, &sep_ho[i].rej_id);
        sep_ho_set.insert(sep_ho[i].t_pos);
    }
    std::fclose(fm);

    auto* tout = g.find("output.weight");
    const size_t rb = tout->bytes / V;
    auto get_wout_row = [&](int vid, std::vector<float>& row) {
        row.resize(N);
        ref::dequant(tout->type, (const uint8_t*)g.map + tout->offset + (size_t)vid * rb, row.data(), N);
    };

    // Notice: in tr_p4 (positions 0..460, the train multi-turn chat region), every token NOT inall_p4_slop_set
    // is a real multi-turn chat background token (user turns, <im_start>, <im_end>, <think>, </think>, and ordinary assistant words)!
    // Weighting those multi-turn chat background tokens strongly in G_base (e.g. w=60.0) AND putting their means
    // in the zero-inner-product constraints forces v_pid . z[t] ~ 0 on ALL ordinary multi-turn chat positions!
    const int K_PHR = 7, K_COL = 6;
    std::vector<double> G_base((size_t)F * F, 0.0);
    auto accum_z_to = [&](std::vector<double>& Gmat, const float* zt, double w) {
        for (int r = 0; r < F; ++r) {
            double zr = w * (double)zt[r];
            for (int c = r; c < F; ++c) Gmat[(size_t)r * F + c] += zr * (double)zt[c];
        }
    };
    for (int t = 0; t < tr_cal.T; ++t) accum_z_to(G_base, &tr_cal.z[(size_t)t * F], 1.0);
    for (int t = 0; t < 480; ++t)      accum_z_to(G_base, &tr_col.z[(size_t)t * F], 4.0);

    std::vector<double> M_chat_bg(F, 0.0), M_chat_imend(F, 0.0), M_cal_bg(F, 0.0);
    std::vector<float> gam_chat_bg(N, 0.0f);
    int n_chat_bg = 0, n_chat_imend = 0;
    for (int t = 0; t < tr_p4.T; ++t) {
        if (all_p4_slop_set.count(t)) continue;
        // Only use the train half of tr_p4 (t < 310 or t >= 460) so heldout_14 context (310..460) is never touched!
        if (t >= 310 && t < 460) continue;
        const float* zt = &tr_p4.z[(size_t)t * F];
        if (t % 2 == 0) accum_z_to(G_base, zt, 2.0);
        int next_tok = tr_p4.toks[t + 1];
        if (next_tok == 248046 || next_tok == 248045 || next_tok == 248068 || next_tok == 248069) {
            for (int k = 0; k < F; ++k) M_chat_imend[k] += (double)zt[k];
            n_chat_imend++;
        } else {
            for (int k = 0; k < F; ++k) M_chat_bg[k] += (double)zt[k];
            const float* gt = &tr_p4.gam[(size_t)t * N];
            for (int j = 0; j < N; ++j) gam_chat_bg[j] += gt[j];
            n_chat_bg++;
        }
    }
    for (int k = 0; k < F; ++k) {
        if (n_chat_bg > 0) M_chat_bg[k] /= n_chat_bg;
        if (n_chat_imend > 0) M_chat_imend[k] /= n_chat_imend;
    }
    if (n_chat_bg > 0) for (int j = 0; j < N; ++j) gam_chat_bg[j] /= (float)n_chat_bg;
    for (int t = 0; t < tr_cal.T; ++t) {
        for (int k = 0; k < F; ++k) M_cal_bg[k] += (double)tr_cal.z[(size_t)t * F + k] / tr_cal.T;
    }

    std::vector<double> M_phr((size_t)K_PHR * F, 0.0), M_col_m((size_t)K_COL * F, 0.0);
    std::vector<int> cnt_phr(K_PHR, 0), cnt_col_m(K_COL, 0);
    int phrase_rej_id[K_PHR] = {};
    std::vector<std::vector<int>> phrase_alts(K_PHR);
    std::vector<float> phrase_gam((size_t)K_PHR * N, 0.0f);

    for (const auto& fr : ftpo) {
        if (fr.split == 0) {
            int pid = fr.phrase_id;
            phrase_rej_id[pid] = fr.rej_id;
            const float* zt = &tr_p4.z[(size_t)fr.t_pos * F];
            const float* gt = &tr_p4.gam[(size_t)fr.t_pos * N];
            for (int k = 0; k < F; ++k) M_phr[(size_t)pid * F + k] += (double)zt[k];
            for (int j = 0; j < N; ++j) phrase_gam[(size_t)pid * N + j] += gt[j];
            cnt_phr[pid]++;
            accum_z_to(G_base, zt, 8.0);
        }
    }
    for (const auto& cr : col) {
        if (cr.split == 0) {
            int b = cr.ban_idx;
            const float* zt = &tr_col.z[(size_t)cr.t_pos * F];
            for (int k = 0; k < F; ++k) M_col_m[(size_t)b * F + k] += (double)zt[k];
            cnt_col_m[b]++;
            accum_z_to(G_base, zt, 40.0);
        }
    }
    for (int p = 0; p < K_PHR; ++p) if (cnt_phr[p] > 0) {
        for (int k = 0; k < F; ++k) M_phr[(size_t)p * F + k] /= cnt_phr[p];
        for (int j = 0; j < N; ++j) phrase_gam[(size_t)p * N + j] /= cnt_phr[p];
    }
    for (int b = 0; b < K_COL; ++b) if (cnt_col_m[b] > 0) {
        for (int k = 0; k < F; ++k) M_col_m[(size_t)b * F + k] /= cnt_col_m[b];
    }

    // Full pass over output.weight for ALL 4 traces: tr_cal (1022), tr_p4 (1022), tr_ho (1022), col (207)
    const int M_cal = tr_cal.T, M_p4 = tr_p4.T, M_ho = tr_ho.T, M_col = (int)col.size();
    const int M_base = M_cal + M_p4 + M_ho + M_col;
    std::vector<float> X_base((size_t)M_base * N);
    for (int t = 0; t < M_cal; ++t) std::memcpy(&X_base[(size_t)t * N], &tr_cal.mix0[(size_t)t * N], N * sizeof(float));
    for (int t = 0; t < M_p4; ++t)  std::memcpy(&X_base[(size_t)(M_cal + t) * N], &tr_p4.mix0[(size_t)t * N], N * sizeof(float));
    for (int t = 0; t < M_ho; ++t)  std::memcpy(&X_base[(size_t)(M_cal + M_p4 + t) * N], &tr_ho.mix0[(size_t)t * N], N * sizeof(float));
    for (int i = 0; i < M_col; ++i) std::memcpy(&X_base[(size_t)(M_cal + M_p4 + M_ho + i) * N], &tr_col.mix0[(size_t)col[i].t_pos * N], N * sizeof(float));

    const int CH = 8192;
    std::vector<float> Wc((size_t)CH * N), L_base((size_t)M_base * CH);
    std::vector<float> mx_b(M_base, -1e30f), se_b(M_base, 0.0f);
    std::vector<int> argmax_b(M_base, -1);
    std::vector<float> y_tgt_b(M_base, 0.0f);
    std::vector<double> w_bar_d(N, 0.0);
    std::vector<float> phrase_train_logits((size_t)K_PHR * V, 0.0f);

    for (int v0 = 0; v0 < V; v0 += CH) {
        int nv = std::min(CH, V - v0);
        ref::dequant(tout->type, (const uint8_t*)g.map + tout->offset + (size_t)v0 * rb, Wc.data(), (size_t)nv * N);
        for (int r = 0; r < nv; ++r) {
            const float* wr = &Wc[(size_t)r * N];
            for (int j = 0; j < N; ++j) w_bar_d[j] += (double)wr[j];
        }
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, M_base, nv, N, 1.0f, X_base.data(), N, Wc.data(), N, 0.0f, L_base.data(), CH);
        for (int m = 0; m < M_base; ++m) {
            const float* l = &L_base[(size_t)m * CH];
            float cm = mx_b[m];
            int best_k = -1;
            for (int k = 0; k < nv; ++k) {
                if (l[k] > cm) { cm = l[k]; best_k = v0 + k; }
            }
            float s = se_b[m] * std::exp(mx_b[m] - cm);
            for (int k = 0; k < nv; ++k) s += std::exp(l[k] - cm);
            if (best_k >= 0) argmax_b[m] = best_k;
            mx_b[m] = cm; se_b[m] = s;
        }
        for (int t = 0; t < M_ho; ++t) {
            int tgt = tr_ho.toks[t + 1];
            if (tgt >= v0 && tgt < v0 + nv) y_tgt_b[M_cal + M_p4 + t] = L_base[(size_t)(M_cal + M_p4 + t) * CH + (tgt - v0)];
        }
        for (const auto& fr : ftpo) {
            if (fr.split == 0) {
                const float* l = &L_base[(size_t)(M_cal + fr.t_pos) * CH];
                float* dst = &phrase_train_logits[(size_t)fr.phrase_id * V + v0];
                for (int k = 0; k < nv; ++k) dst[k] += l[k] / (float)cnt_phr[fr.phrase_id];
            }
        }
    }
    std::vector<float> lse_b(M_base);
    for (int m = 0; m < M_base; ++m) lse_b[m] = mx_b[m] + std::log(se_b[m]);
    std::vector<float> w_bar(N);
    for (int j = 0; j < N; ++j) w_bar[j] = (float)(w_bar_d[j] / V);

    // Exclude banned words, special tokens, punctuation, and ultra-common function words (like 279 " the", 760 "Let") from phrase_alts
    // so boosting phrase_alts[pid] NEVER perturbs ordinary chat tokens when v_z < 0!
    // Explicit clean natural non-common alternatives per phrase (from train_logits / grammar):
    phrase_alts[0] = {1151, 3839};  // "end to": " start" (1151), " End" (3839)
    phrase_alts[1] = {1204, 1332};  // "exactly": " how" (1204), " where" (1332)
    phrase_alts[2] = {1204, 1332};  // "here's": " how" (1204), " where" (1332)
    phrase_alts[3] = {383, 4237};   // "rather": " on" (383), " instead" (4237)
    phrase_alts[4] = {1238, 1518};  // "right": " after" (1238), " before" (1518)
    phrase_alts[5] = {4087, 3255};  // "short": " answer" (4087), " story" (3255)
    phrase_alts[6] = {12125, 2019}; // "what's": " happening" (12125), " going" (2019)

    int ban_ids_arr[6] = {809, 1056, 1092, 2243, 3023, 3404};
    const char* ban_names[6] = {" end", " than", " what", " version", " away", " actually"};
    auto find_ban_idx = [&](int rid) {
        for (int i = 0; i < 6; ++i) if (ban_ids_arr[i] == rid) return i;
        return 0;
    };

    // Solve per-phrase v_pid in R^640 with exact orthogonality to:
    // - M_col_m[bidx] (collateral mean for that word)
    // - M_cal_bg (calibration mean)
    // - tr_p4.z[705] (SFT anchor for pid=0,3)
    std::vector<double> V_cls((size_t)K_PHR * F, 0.0);
    for (int pid = 0; pid < K_PHR; ++pid) {
        int bidx = find_ban_idx(phrase_rej_id[pid]);
        std::vector<double> G_b = G_base;
        double trG_base = 0.0;
        for (int k = 0; k < F; ++k) trG_base += G_base[(size_t)k * F + k];
        double ridge = 1.20 * (trG_base / F);
        double col_w = (pid == 0 || pid == 3 || pid == 4 || pid == 6) ? 12000.0 : 6500.0;
        for (const auto& cr : col) {
            if (cr.split == 0) {
                double w_cr = (cr.ban_idx == bidx) ? col_w : 400.0;
                accum_z_to(G_b, &tr_col.z[(size_t)cr.t_pos * F], w_cr);
            }
        }
        for (int r = 0; r < F; ++r) {
            G_b[(size_t)r * F + r] += ridge;
            for (int c = 0; c < r; ++c) G_b[(size_t)r * F + c] = G_b[(size_t)c * F + r];
        }

        std::vector<std::vector<double>> c_rows;
        c_rows.push_back(std::vector<double>(&M_phr[(size_t)pid * F], &M_phr[(size_t)(pid + 1) * F]));
        if (pid == 1) c_rows.push_back(std::vector<double>(&M_phr[2 * F], &M_phr[3 * F]));
        if (pid == 2) c_rows.push_back(std::vector<double>(&M_phr[1 * F], &M_phr[2 * F]));
        c_rows.push_back(std::vector<double>(&M_col_m[(size_t)bidx * F], &M_col_m[(size_t)(bidx + 1) * F]));
        c_rows.push_back(M_cal_bg);
        if (pid == 0 || pid == 3) {
            c_rows.push_back(std::vector<double>(&tr_p4.z[705 * F], &tr_p4.z[706 * F]));
        }

        int ks = (int)c_rows.size();
        std::vector<double> Q_sub((size_t)ks * F, 0.0);
        for (int i = 0; i < ks; ++i) for (int k = 0; k < F; ++k) Q_sub[(size_t)i * F + k] = c_rows[i][k];
        char uplo = 'L';
        int n_lap = F, nrhs = ks, lda = F, ldb = F, info = 0;
        dposv_(&uplo, &n_lap, &nrhs, G_b.data(), &lda, Q_sub.data(), &ldb, &info);

        std::vector<double> S_sub((size_t)ks * ks, 0.0);
        for (int i = 0; i < ks; ++i) {
            for (int j = 0; j < ks; ++j) {
                double s = 0.0;
                for (int k = 0; k < F; ++k) s += c_rows[i][k] * Q_sub[(size_t)j * F + k];
                S_sub[(size_t)i * ks + j] = s;
            }
            S_sub[(size_t)i * ks + i] += 1e-6;
        }
        std::vector<double> rhs(ks, 0.0);
        rhs[0] = 1.0;
        int nr = 1, info_s = 0;
        dposv_(&uplo, &ks, &nr, S_sub.data(), &ks, rhs.data(), &ks, &info_s);
        for (int i = 0; i < ks; ++i) {
            for (int k = 0; k < F; ++k) V_cls[(size_t)pid * F + k] += rhs[i] * Q_sub[(size_t)i * F + k];
        }
    }

    // Output-space u_pid in R^2560:
    // Constrain u_pid against:
    //   0: a_rej      -> -delta_rej
    //   1: a_cho      -> +delta_cho (natural alternatives)
    //   2: a_im_end   (248046) -> -0.15 * delta_rej (slight suppression of <|im_end|> so it NEVER rises!)
    //   3: a_im_start (248045) ->  0.0
    //   4: a_think    (248068) ->  0.0
    //   5: a_ethink   (248069) ->  0.0
    //   6: a_eos      (248044) ->  0.0
    int spec_ids[11] = {248046, 248045, 248068, 248069, 248044, 279, 760, 11, 13, 198, 3437};
    std::vector<float> w_spec[11];
    for (int s = 0; s < 11; ++s) get_wout_row(spec_ids[s], w_spec[s]);

    double phrase_scales[7] = {3.7, 5.4, 5.4, 4.2, 4.1, 5.5, 3.5};
    for (double mult : {0.88}) {
        std::vector<float> dW((size_t)N * F, 0.0f);
        for (int pid = 0; pid < K_PHR; ++pid) {
            double delta_rej = mult * phrase_scales[pid];
            std::vector<float> w_r, wc1, wc2;
            get_wout_row(phrase_rej_id[pid], w_r);
            get_wout_row(phrase_alts[pid][0], wc1);
            get_wout_row(phrase_alts[pid][1], wc2);
            const float* gm = &phrase_gam[(size_t)pid * N];
            double d_c1 = (pid == 0) ? (4.50 * delta_rej) : ((pid == 3) ? (0.65 * delta_rej) : (0.16 * delta_rej));
            double d_c2 = (pid == 0) ? (3.00 * delta_rej) : ((pid == 3) ? (0.40 * delta_rej) : (0.10 * delta_rej));
            double d_imend = (pid == 0) ? (-0.28 * delta_rej) : ((pid == 3) ? (-0.26 * delta_rej) : (-0.02 * delta_rej));
            const int NU = 15;
            std::vector<std::vector<double>> A(NU, std::vector<double>(N));
            for (int j = 0; j < N; ++j) {
                A[0][j] = (double)gm[j] * (double)w_r[j];
                A[1][j] = (double)gm[j] * (double)wc1[j];
                A[2][j] = (double)gm[j] * (double)wc2[j];
                A[3][j] = (double)gm[j] * (double)w_bar[j];
                for (int s = 0; s < 11; ++s) {
                    A[4 + s][j] = (double)gm[j] * (double)w_spec[s][j];
                }
            }
            std::vector<double> G_u(NU * NU, 0.0);
            for (int i = 0; i < NU; ++i) {
                for (int k = 0; k < NU; ++k) {
                    double dot = 0.0;
                    for (int j = 0; j < N; ++j) dot += A[i][j] * A[k][j];
                    G_u[i * NU + k] = dot;
                }
                G_u[i * NU + i] += 1e-6;
            }
            std::vector<double> rhs_u = {-delta_rej, +d_c1, +d_c2, 0.0, d_imend, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0};
            char uplo = 'L';
            int nu_lap = NU, nr = 1, info_u = 0;
            dposv_(&uplo, &nu_lap, &nr, G_u.data(), &nu_lap, rhs_u.data(), &nu_lap, &info_u);

            for (int j = 0; j < N; ++j) {
                double u_j = 0.0;
                for (int i = 0; i < NU; ++i) u_j += rhs_u[i] * A[i][j];
                for (int k = 0; k < F; ++k) dW[(size_t)j * F + k] += (float)u_j * (float)V_cls[(size_t)pid * F + k];
            }
        }

        std::vector<float> w_target((size_t)N * F);
        for (size_t i = 0; i < w_target.size(); ++i) w_target[i] = wb[i] + dW[i];
        std::vector<uint8_t> raw_q8((size_t)N * F / 32 * 34);
        quantize_q8_0_monotonic(raw_sd47_orig, w_target.data(), w_target.size(), raw_q8.data());
        std::vector<float> w_q8_deq((size_t)N * F), dq((size_t)N * F);
        FILE* fq = std::fopen("/tmp/blk47_v3_q8.bin", "wb");
        std::fwrite(raw_q8.data(), 1, raw_q8.size(), fq);
        std::fclose(fq);
        ref::dequant(ref::Q8_0, raw_q8.data(), w_q8_deq.data(), w_q8_deq.size());
        double dn = 0, bn = 0;
        for (size_t i = 0; i < dq.size(); ++i) {
            dq[i] = w_q8_deq[i] - wb[i];
            dn += (double)dq[i] * dq[i];
            bn += (double)wb[i] * wb[i];
        }

        TraceData* trs[4] = {&tr_cal, &tr_p4, &tr_ho, &tr_col};
        std::vector<float> mix_new[4];
        for (int ti = 0; ti < 4; ++ti) {
            int T = trs[ti]->T;
            std::vector<float> d_t((size_t)T * N);
            cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, T, N, F, 1.0f, trs[ti]->z.data(), F, dq.data(), F, 0.0f, d_t.data(), N);
            mix_new[ti].resize((size_t)T * N);
            compute_mix_batch(T, trs[ti]->R.data(), trs[ti]->inj.data(), d_t.data(), wn, wdn, wup, mix_new[ti].data(), nullptr);
        }

        std::vector<float> X_new((size_t)M_base * N);
        for (int t = 0; t < M_cal; ++t) std::memcpy(&X_new[(size_t)t * N], &mix_new[0][(size_t)t * N], N * sizeof(float));
        for (int t = 0; t < M_p4; ++t)  std::memcpy(&X_new[(size_t)(M_cal + t) * N], &mix_new[1][(size_t)t * N], N * sizeof(float));
        for (int t = 0; t < M_ho; ++t)  std::memcpy(&X_new[(size_t)(M_cal + M_p4 + t) * N], &mix_new[2][(size_t)t * N], N * sizeof(float));
        for (int i = 0; i < M_col; ++i) std::memcpy(&X_new[(size_t)(M_cal + M_p4 + M_ho + i) * N], &mix_new[3][(size_t)col[i].t_pos * N], N * sizeof(float));

        std::vector<float> L_new((size_t)M_base * CH), mx_n(M_base, -1e30f), se_n(M_base, 0.0f);
        std::vector<int> argmax_n(M_base, -1);
        std::vector<float> y_tgt_n(M_base, 0.0f);
        std::vector<double> dot_sum(M_base, 0.0);

        for (int v0 = 0; v0 < V; v0 += CH) {
            int nv = std::min(CH, V - v0);
            ref::dequant(tout->type, (const uint8_t*)g.map + tout->offset + (size_t)v0 * rb, Wc.data(), (size_t)nv * N);
            cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, M_base, nv, N, 1.0f, X_base.data(), N, Wc.data(), N, 0.0f, L_base.data(), CH);
            cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, M_base, nv, N, 1.0f, X_new.data(), N, Wc.data(), N, 0.0f, L_new.data(), CH);
            for (int m = 0; m < M_base; ++m) {
                const float* lb = &L_base[(size_t)m * CH];
                const float* ln = &L_new[(size_t)m * CH];
                float cm = mx_n[m];
                int best_k = -1;
                for (int k = 0; k < nv; ++k) {
                    if (ln[k] > cm) { cm = ln[k]; best_k = v0 + k; }
                }
                float s = se_n[m] * std::exp(mx_n[m] - cm);
                double ds = 0.0;
                float lseb = lse_b[m];
                for (int k = 0; k < nv; ++k) {
                    s += std::exp(ln[k] - cm);
                    float pb = std::exp(lb[k] - lseb);
                    ds += (double)pb * ((double)lb[k] - (double)ln[k]);
                }
                if (best_k >= 0) argmax_n[m] = best_k;
                mx_n[m] = cm; se_n[m] = s;
                dot_sum[m] += ds;
            }
            for (int t = 0; t < M_ho; ++t) {
                int tgt = tr_ho.toks[t + 1];
                if (tgt >= v0 && tgt < v0 + nv) y_tgt_n[M_cal + M_p4 + t] = L_new[(size_t)(M_cal + M_p4 + t) * CH + (tgt - v0)];
            }
        }
        std::vector<float> lse_n(M_base);
        for (int m = 0; m < M_base; ++m) lse_n[m] = mx_n[m] + std::log(se_n[m]);

        // Check bars (f) and (g) on held-out chat doc (tr_ho, M_ho = 1022):
        int flips_to_spec_ho = 0, flips_to_spec_p4 = 0, flips_to_spec_slop21 = 0;
        auto is_spec = [](int tok) { return tok == 248046 || tok == 248045 || tok == 248068 || tok == 248069 || tok == 248044; };
        double ord_dlp_ho = 0.0, ho_kld_sum = 0.0, nonslop_kld_sum = 0.0;
        int n_ord_ho = 0, flips_nonslop_ho = 0;
        for (int t = 0; t < M_ho; ++t) {
            int m = M_cal + M_p4 + t;
            int tgt = tr_ho.toks[t + 1];
            double kld_t = dot_sum[m] - (double)lse_b[m] + (double)lse_n[m];
            ho_kld_sum += kld_t;
            if (!sep_ho_set.count(t)) nonslop_kld_sum += kld_t;
            if (!is_spec(tgt) && !is_spec(argmax_b[m]) && is_spec(argmax_n[m])) {
                flips_to_spec_ho++;
                if (!sep_ho_set.count(t)) {
                    flips_nonslop_ho++;
                    std::printf("    [NONSLOP FLIP t=%d] tgt=%d argmax_b=%d argmax_n=%d\n", t, tgt, argmax_b[m], argmax_n[m]);
                }
            }
            if (!sep_ho_set.count(t) && tgt != 248046) {
                double dlp = ((double)y_tgt_n[m] - lse_n[m]) - ((double)y_tgt_b[m] - lse_b[m]);
                ord_dlp_ho += dlp;
                n_ord_ho++;
            }
        }
        for (int t = 0; t < M_p4; ++t) {
            int m = M_cal + t;
            int tgt = tr_p4.toks[t + 1];
            if (!is_spec(tgt) && !is_spec(argmax_b[m]) && is_spec(argmax_n[m])) {
                flips_to_spec_p4++;
                std::printf("    [SFT FLIP t=%d is_slop=%d] tgt=%d argmax_b=%d argmax_n=%d\n", t, (int)all_p4_slop_set.count(t), tgt, argmax_b[m], argmax_n[m]);
            }
        }

        // Check FTPO Train (28), Heldout_14, and Separate-Doc Heldout_21 (plus p(<|im_end|>) at slop21!)
        double p_tr_b = 0, p_tr_n = 0; int n_tr = 0;
        double p_ho_b = 0, p_ho_n = 0; int n_ho = 0;
        for (const auto& fr : ftpo) {
            int m = M_cal + fr.t_pos;
            std::vector<float> wr; get_wout_row(fr.rej_id, wr);
            float yr_b = 0, yr_n = 0;
            for (int j = 0; j < N; ++j) { yr_b += wr[j] * X_base[(size_t)m * N + j]; yr_n += wr[j] * X_new[(size_t)m * N + j]; }
            double pr_b = std::exp((double)yr_b - lse_b[m]);
            double pr_n = std::exp((double)yr_n - lse_n[m]);
            if (fr.split == 0) { p_tr_b += pr_b; p_tr_n += pr_n; n_tr++; }
            else               { p_ho_b += pr_b; p_ho_n += pr_n; n_ho++; }
        }
        double p_sep_b = 0, p_sep_n = 0, pim_sep_b = 0, pim_sep_n = 0;
        for (int i = 0; i < n_sep; ++i) {
            int m = M_cal + M_p4 + sep_ho[i].t_pos;
            if (false) {
                std::vector<std::pair<float, int>> top5;
                for (int v0 = 0; v0 < V; v0 += CH) {
                    int nv = std::min(CH, V - v0);
                    ref::dequant(tout->type, (const uint8_t*)g.map + tout->offset + (size_t)v0 * rb, Wc.data(), (size_t)nv * N);
                    for (int k = 0; k < nv; ++k) {
                        float yv = 0;
                        const float* wr = &Wc[(size_t)k * N];
                        for (int j = 0; j < N; ++j) yv += wr[j] * X_base[(size_t)m * N + j];
                        if (yv > 12.0f) top5.push_back({yv, v0 + k});
                    }
                }
                std::sort(top5.rbegin(), top5.rend());
                std::printf("  slop21[%2d] pid=%d t=%3d rej=%4d: ", i, sep_ho[i].phrase_id, sep_ho[i].t_pos, sep_ho[i].rej_id);
                for (int k = 0; k < std::min<int>(5, top5.size()); ++k) std::printf("%d(%.2f) ", top5[k].second, top5[k].first);
                std::printf("\n");
            }
            if (!is_spec(argmax_b[m]) && is_spec(argmax_n[m])) flips_to_spec_slop21++;
            std::vector<float> wr; get_wout_row(sep_ho[i].rej_id, wr);
            float yr_b = 0, yr_n = 0, yim_b = 0, yim_n = 0;
            for (int j = 0; j < N; ++j) {
                yr_b += wr[j] * X_base[(size_t)m * N + j];
                yr_n += wr[j] * X_new[(size_t)m * N + j];
                yim_b += w_spec[0][j] * X_base[(size_t)m * N + j];
                yim_n += w_spec[0][j] * X_new[(size_t)m * N + j];
            }
            p_sep_b += std::exp((double)yr_b - lse_b[m]);
            p_sep_n += std::exp((double)yr_n - lse_n[m]);
            pim_sep_b += std::exp((double)yim_b - lse_b[m]);
            pim_sep_n += std::exp((double)yim_n - lse_n[m]);
        }

        // Per-word Held-Out Collateral breakdown (n=142)
        double w_sum[6] = {}, w_min[6] = {1e9, 1e9, 1e9, 1e9, 1e9, 1e9};
        int w_cnt[6] = {}, w_loss10[6] = {}, tot_loss10 = 0, nc_ho = 0;
        double dlp_ho_all = 0;
        for (int i = 0; i < M_col; ++i) {
            if (col[i].split == 0) continue;
            int m = M_cal + M_p4 + M_ho + i;
            std::vector<float> wt; get_wout_row(col[i].target_id, wt);
            float yt_b = 0, yt_n = 0;
            for (int j = 0; j < N; ++j) { yt_b += wt[j] * X_base[(size_t)m * N + j]; yt_n += wt[j] * X_new[(size_t)m * N + j]; }
            double dlp = ((double)yt_n - lse_n[m]) - ((double)yt_b - lse_b[m]);
            int b = col[i].ban_idx;
            w_sum[b] += dlp; w_cnt[b]++; w_min[b] = std::min(w_min[b], dlp);
            if (dlp < -0.10) { w_loss10[b]++; tot_loss10++; }
            dlp_ho_all += dlp; nc_ho++;
        }

        std::printf("\n=== v3.2 mult=%.2f (||dW||/||W||=%.4f) ===\n", mult, std::sqrt(dn / bn));
        std::printf("  FTPO TRAIN       (n=%d): %.5f -> %.5f (-%.2f%%)\n", n_tr, p_tr_b / n_tr, p_tr_n / n_tr, (1.0 - p_tr_n / p_tr_b) * 100.0);
        std::printf("  FTPO HELDOUT_DOC (n=%d): %.5f -> %.5f (-%.2f%%) | p(<|im_end|>): %.5f -> %.5f\n", n_sep, p_sep_b / n_sep, p_sep_n / n_sep, (1.0 - p_sep_n / p_sep_b) * 100.0, pim_sep_b / n_sep, pim_sep_n / n_sep);
        std::printf("  FTPO HELDOUT_14  (n=%d): %.5f -> %.5f (-%.2f%%)\n", n_ho, p_ho_b / n_ho, p_ho_n / n_ho, (1.0 - p_ho_n / p_ho_b) * 100.0);
        std::printf("  BAR (f) Special-Token Flips: ho_nonslop=%d (bar 0) | slop21=%d (bar 0) | ho_tot=%d | sft_doc=%d\n", flips_nonslop_ho, flips_to_spec_slop21, flips_to_spec_ho, flips_to_spec_p4);
        std::printf("  BAR (g) Held-Out Chat Doc (n=%d): ord_dlp=%+.5f (bar >= -0.01) | nonslop_KLD=%.5f | whole_KLD=%.5f\n", n_ord_ho, ord_dlp_ho / n_ord_ho, nonslop_kld_sum / (M_ho - n_sep), ho_kld_sum / M_ho);
        std::printf("  HELDOUT COLLATERAL overall (n=%d): mean=%+.5f | loss>0.1: %d/%d (%.1f%%, bar <= 10.0%%)\n",
                    nc_ho, dlp_ho_all / nc_ho, tot_loss10, nc_ho, 100.0 * tot_loss10 / nc_ho);
        for (int b = 0; b < 6; ++b) {
            std::printf("    word %-10s (n=%2d): mean=%+.5f (bar >= -0.05) | worst=%+.4f | <-0.1: %d\n",
                        ban_names[b], w_cnt[b], w_sum[b] / std::max(1, w_cnt[b]), w_min[b], w_loss10[b]);
        }
    }
    return 0;
}