#!/usr/bin/env python3
"""Patch vendor/strata/metal for:
1. Dynamic n_expert (K=256 / 320 / 224 / 512) in router.metal, moe.metal, prompt.metal,
   stream.h, stream.cpp, window.cpp, prompt_chain.inc, and route_parity.cpp (with K=256 + MoCE parity test).
2. Zero-copy file-backed dense tensor wrapping (cx.wrap(f.map, max_dense_end, false)) in
   ExpertStream::load_dense with dynamically computed locked/skipped/unpinned byte totals.
3. Support for dense-only <model>-metal.gguf paired with STRX <model>-metal.gguf.experts.
4. Race-free MoCE cache-prior reranking (STRATA_METAL_MOCE_BONUS) using a per-lane 16-bit register
   mask (scaled_mask + simd_broadcast) in router.metal / moe.metal, plus top-3 unperturbed protection
   and smap residency consistency in ExpertStream::ahead.
5. Transient PromptGemm scratch buffer mapped on demand, skipped memset on fresh MAP_ANON pages,
   unmapped via munmap + cx_->free after prefill with measured proc_mem() footprint logging,
   router GEMM over n_exp_ rows, and multi-turn hot-slot protection across prompts.
6. Standalone zero-bloat strata_window_cli (with --logits-out / --logits-ref full-vocab FP64 KLD
   & PPL evaluator) and convert_k256_q2_0 targets in metal/CMakeLists.txt.
"""

import subprocess
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent / "vendor" / "strata"


def reset_git() -> None:
    subprocess.run(["git", "-C", str(ROOT), "checkout", "--", "."], check=True)


def patch_mtl_rt() -> None:
    h_path = ROOT / "metal" / "src" / "mtl_rt.h"
    mm_path = ROOT / "metal" / "src" / "mtl_rt.mm"

    h_txt = h_path.read_text(encoding="utf-8")
    h_txt = h_txt.replace(
        "Buffer wrap(void* page_aligned, size_t bytes);  // no-copy wrap (page-aligned, size multiple of page)",
        "Buffer wrap(void* page_aligned, size_t bytes, bool count_allocated = true);  // no-copy wrap (page-aligned, size multiple of page)",
    )
    h_path.write_text(h_txt, encoding="utf-8")

    mm_txt = mm_path.read_text(encoding="utf-8")
    old_wrap = """Buffer Context::wrap(void* p, size_t bytes) {
    id<MTLBuffer> b = [impl_->dev newBufferWithBytesNoCopy:p length:bytes options:MTLResourceStorageModeShared deallocator:nil];
    if (!b) {
        std::fprintf(stderr, "mtl: wrap %zu failed\\n", bytes);
        std::exit(4);
    }
    allocated_ += bytes;"""
    new_wrap = """Buffer Context::wrap(void* p, size_t bytes, bool count_allocated) {
    id<MTLBuffer> b = [impl_->dev newBufferWithBytesNoCopy:p length:bytes options:MTLResourceStorageModeShared deallocator:nil];
    if (!b) {
        std::fprintf(stderr, "mtl: wrap %zu failed\\n", bytes);
        std::exit(4);
    }
    if (count_allocated) allocated_ += bytes;"""
    assert old_wrap in mm_txt, "wrap target not found in mtl_rt.mm"
    mm_path.write_text(mm_txt.replace(old_wrap, new_wrap, 1), encoding="utf-8")


def patch_kernels() -> None:
    # 1. router.metal: dynamic a.n_expert + race-free MoCE cache-prior reranking via a.pad, smap & scaled_mask
    r_path = ROOT / "metal" / "kernels" / "router.metal"
    r_txt = r_path.read_text(encoding="utf-8")
    old_route_body = """static inline void route_body(device const float* logits, device int* ids, device float* w, device int* plan,
                              constant MoeArgs& a, ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp,
                              threadgroup int* s_ent) {
    if (sg < a.n_tok) {
        device const float* l = logits + sg * 512;
        float v[16];
        for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i];
        float mx = -INFINITY;
        for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
        for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
        float sum = 0.0f;
        for (int i = 0; i < 16; ++i) {
            v[i] = exp(v[i] - mx);
            sum += v[i];
        }
        for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
        const float rcp = 1.0f / sum;
        for (int i = 0; i < 16; ++i) {
            v[i] *= rcp;
            if (isnan(v[i])) v[i] = -FLT_MAX;
        }
        float selected = 0.0f, ssum = 0.0f;
        for (int rank = 0; rank < 10; ++rank) {
            float best = v[0];
            int ex = lane;
            for (int i = 1; i < 16; ++i)
                if (v[i] > best) { best = v[i]; ex = lane + 32 * i; }
            for (ushort m = 16; m > 0; m >>= 1) {
                const float ob = simd_shuffle_xor(best, m);
                const int oe = simd_shuffle_xor(ex, m);
                if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; }
            }
            if ((ex & 31) == lane) {
                v[ex / 32] = -INFINITY;
                ids[sg * 10 + rank] = ex;
                ssum += best;
            }
            if (rank == lane) selected = best;
        }"""
    new_route_body = """#ifndef STRATA_MOCE_FC_METAL
#define STRATA_MOCE_FC_METAL
constant int kMoceProtectFc [[function_constant(24)]];
constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3;
constant uint kMoceScaleR1BitsFc [[function_constant(25)]];
constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type<float>(kMoceScaleR1BitsFc) : 1.0f;
#endif
static inline void route_body(device const float* logits, device int* ids, device float* w, device int* plan,
                              constant MoeArgs& a, device const uint* smap, ushort tid, ushort nth, ushort sg, ushort lane,
                              threadgroup int* s_grp, threadgroup int* s_ent) {
    if (sg < a.n_tok) {
        device const float* l = logits + sg * 512;
        float v[16];
        for (int i = 0; i < 16; ++i) {
            const int idx = lane + 32 * i;
            v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY;
        }
        float mx = -INFINITY;
        for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
        for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
        float sum = 0.0f;
        for (int i = 0; i < 16; ++i) {
            v[i] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f;
            sum += v[i];
        }
        for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
        const float rcp = 1.0f / sum;
        for (int i = 0; i < 16; ++i) {
            if (lane + 32 * i >= int(a.n_expert)) {
                v[i] = -INFINITY;
            } else {
                v[i] *= rcp;
                if (isnan(v[i])) v[i] = -FLT_MAX;
            }
        }
        const float moce_scale = a.pad != 0u ? as_type<float>(a.pad) : 1.0f;
        const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f;
        const int r_prot = kMoceProtect;
        const float r1_scale = kMoceScaleR1;
        const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f;
        const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale;
        ushort scaled_mask = 0;
        float selected = 0.0f, ssum = 0.0f;
        for (int rank = 0; rank < 10; ++rank) {
            if (r1_scale > 1.0f && r_prot > 1) {
                if (rank == 1 && smap != nullptr) {
                    for (int i = 0; i < 16; ++i) {
                        const int idx = lane + 32 * i;
                        if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) {
                            v[i] *= r1_scale;
                            scaled_mask |= ushort(1u << i);
                        }
                    }
                } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) {
                    for (int i = 0; i < 16; ++i) {
                        if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) {
                            v[i] *= r_step_scale;
                        }
                    }
                }
            } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) {
                for (int i = 0; i < 16; ++i) {
                    const int idx = lane + 32 * i;
                    if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) {
                        v[i] *= moce_scale;
                        scaled_mask |= ushort(1u << i);
                    }
                }
            }
            float best = v[0];
            int ex = lane;
            for (int i = 1; i < 16; ++i)
                if (v[i] > best) { best = v[i]; ex = lane + 32 * i; }
            for (ushort m = 16; m > 0; m >>= 1) {
                const float ob = simd_shuffle_xor(best, m);
                const int oe = simd_shuffle_xor(ex, m);
                if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; }
            }
            float true_prob = best;
            const int min_scaled_rank = (r1_scale > 1.0f && r_prot > 1) ? 1 : r_prot;
            if (rank >= min_scaled_rank && (moce_scale > 1.0f || r1_scale > 1.0f) && best > 0.0f) {
                const ushort win_lane = ushort(ex & 31);
                const ushort win_i = ushort(ex >> 5);
                const ushort win_mask = simd_broadcast(scaled_mask, win_lane);
                if ((win_mask & ushort(1u << win_i)) != 0u) {
                    true_prob = best * ((r1_scale > 1.0f && rank < r_prot) ? rcp_r1 : rcp_moce);
                }
            }
            if ((ex & 31) == lane) {
                v[ex / 32] = -INFINITY;
                ids[sg * 10 + rank] = ex;
                ssum += true_prob;
            }
            if (rank == lane) selected = true_prob;
        }"""
    assert old_route_body in r_txt, "route_body not found in router.metal"
    r_txt = r_txt.replace(old_route_body, new_route_body, 1)
    r_txt = r_txt.replace(
        "device atomic_uint* counter [[buffer(9)]],",
        "device atomic_uint* counter [[buffer(9)]], device const uint* smap [[buffer(10)]],",
    )
    r_txt = r_txt.replace(
        "route_body(logits, ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent);",
        "route_body(logits, ids, w, plan, a, smap, tid, nth, sg, lane, s_grp, s_ent);",
    )
    r_path.write_text(r_txt, encoding="utf-8")

    # 2. moe.metal: dynamic a.n_expert + race-free MoCE cache-prior reranking in moe_route
    m_path = ROOT / "metal" / "kernels" / "moe.metal"
    m_txt = m_path.read_text(encoding="utf-8")
    old_moe_route = """kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]],
                      device float* w [[buffer(2)]], device int* plan [[buffer(3)]], constant MoeArgs& a [[buffer(4)]],
                      ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]],
                      ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
    threadgroup int s_grp[32], s_ent[32];
    if (sg < a.n_tok) {
        device const float* l = logits + sg * 512;
        float v[16];
        for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i];
        float mx = -INFINITY;
        for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
        for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
        float sum = 0.0f;
        for (int i = 0; i < 16; ++i) {
            v[i] = exp(v[i] - mx);
            sum += v[i];
        }
        for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
        const float rcp = 1.0f / sum;
        for (int i = 0; i < 16; ++i) {
            v[i] *= rcp;
            if (isnan(v[i])) v[i] = -FLT_MAX;
        }
        float selected = 0.0f, ssum = 0.0f;
        for (int rank = 0; rank < 10; ++rank) {
            float best = v[0];
            int ex = lane;
            for (int i = 1; i < 16; ++i)
                if (v[i] > best) { best = v[i]; ex = lane + 32 * i; }
            for (ushort m = 16; m > 0; m >>= 1) {
                const float ob = simd_shuffle_xor(best, m);
                const int oe = simd_shuffle_xor(ex, m);
                if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; }
            }
            if ((ex & 31) == lane) {
                v[ex / 32] = -INFINITY;
                ids[sg * 10 + rank] = ex;
                ssum += best;
            }
            if (rank == lane) selected = best;
        }"""
    new_moe_route = """#ifndef STRATA_MOCE_FC_METAL
#define STRATA_MOCE_FC_METAL
constant int kMoceProtectFc [[function_constant(24)]];
constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3;
constant uint kMoceScaleR1BitsFc [[function_constant(25)]];
constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type<float>(kMoceScaleR1BitsFc) : 1.0f;
#endif
kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]],
                      device float* w [[buffer(2)]], device int* plan [[buffer(3)]], constant MoeArgs& a [[buffer(4)]],
                      device const uint* smap [[buffer(5)]],
                      ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]],
                      ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
    threadgroup int s_grp[32], s_ent[32];
    if (sg < a.n_tok) {
        device const float* l = logits + sg * 512;
        float v[16];
        for (int i = 0; i < 16; ++i) {
            const int idx = lane + 32 * i;
            v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY;
        }
        float mx = -INFINITY;
        for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
        for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
        float sum = 0.0f;
        for (int i = 0; i < 16; ++i) {
            v[i] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f;
            sum += v[i];
        }
        for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
        const float rcp = 1.0f / sum;
        for (int i = 0; i < 16; ++i) {
            if (lane + 32 * i >= int(a.n_expert)) {
                v[i] = -INFINITY;
            } else {
                v[i] *= rcp;
                if (isnan(v[i])) v[i] = -FLT_MAX;
            }
        }
        const float moce_scale = a.pad != 0u ? as_type<float>(a.pad) : 1.0f;
        const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f;
        const int r_prot = kMoceProtect;
        const float r1_scale = kMoceScaleR1;
        const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f;
        const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale;
        ushort scaled_mask = 0;
        float selected = 0.0f, ssum = 0.0f;
        for (int rank = 0; rank < 10; ++rank) {
            if (r1_scale > 1.0f && r_prot > 1) {
                if (rank == 1 && smap != nullptr) {
                    for (int i = 0; i < 16; ++i) {
                        const int idx = lane + 32 * i;
                        if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) {
                            v[i] *= r1_scale;
                            scaled_mask |= ushort(1u << i);
                        }
                    }
                } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) {
                    for (int i = 0; i < 16; ++i) {
                        if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) {
                            v[i] *= r_step_scale;
                        }
                    }
                }
            } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) {
                for (int i = 0; i < 16; ++i) {
                    const int idx = lane + 32 * i;
                    if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) {
                        v[i] *= moce_scale;
                        scaled_mask |= ushort(1u << i);
                    }
                }
            }
            float best = v[0];
            int ex = lane;
            for (int i = 1; i < 16; ++i)
                if (v[i] > best) { best = v[i]; ex = lane + 32 * i; }
            for (ushort m = 16; m > 0; m >>= 1) {
                const float ob = simd_shuffle_xor(best, m);
                const int oe = simd_shuffle_xor(ex, m);
                if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; }
            }
            float true_prob = best;
            const int min_scaled_rank = (r1_scale > 1.0f && r_prot > 1) ? 1 : r_prot;
            if (rank >= min_scaled_rank && (moce_scale > 1.0f || r1_scale > 1.0f) && best > 0.0f) {
                const ushort win_lane = ushort(ex & 31);
                const ushort win_i = ushort(ex >> 5);
                const ushort win_mask = simd_broadcast(scaled_mask, win_lane);
                if ((win_mask & ushort(1u << win_i)) != 0u) {
                    true_prob = best * ((r1_scale > 1.0f && rank < r_prot) ? rcp_r1 : rcp_moce);
                }
            }
            if ((ex & 31) == lane) {
                v[ex / 32] = -INFINITY;
                ids[sg * 10 + rank] = ex;
                ssum += true_prob;
            }
            if (rank == lane) selected = true_prob;
        }"""
    assert old_moe_route in m_txt, "moe_route target not found in moe.metal"
    m_path.write_text(m_txt.replace(old_moe_route, new_moe_route, 1), encoding="utf-8")

    # 3. prompt.metal: route_topk unpacks n_tok = n_tok_packed & 0xFFFF, n_exp = (n_tok_packed >> 16) ? (n_tok_packed >> 16) : 512,
    #    and indexes row t at logits + ulong(t) * ulong(n_exp) to match gemm(..., T, n_exp_, N)
    p_path = ROOT / "metal" / "kernels" / "prompt.metal"
    p_txt = p_path.read_text(encoding="utf-8")
    old_ptopk = """kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]],
                       device float* w [[buffer(2)]], constant uint& n_tok [[buffer(3)]],
                       uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
                       ushort lane [[thread_index_in_simdgroup]]) {
    const uint t = tg.x * 8 + sg;
    if (t >= n_tok) return;
    device const float* l = logits + ulong(t) * 512;
    float v[16];
    for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i];
    float mx = -INFINITY;
    for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
    for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
    float sum = 0.0f;
    for (int i = 0; i < 16; ++i) {
        v[i] = exp(v[i] - mx);
        sum += v[i];
    }
    for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
    const float rcp = 1.0f / sum;
    for (int i = 0; i < 16; ++i) {
        v[i] *= rcp;
        if (isnan(v[i])) v[i] = -FLT_MAX;
    }"""
    new_ptopk = """kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]],
                       device float* w [[buffer(2)]], constant uint& n_tok_packed [[buffer(3)]],
                       uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
                       ushort lane [[thread_index_in_simdgroup]]) {
    const uint n_tok = n_tok_packed & 0xFFFFu;
    const int n_exp = (n_tok_packed >> 16) ? int(n_tok_packed >> 16) : 512;
    const uint t = tg.x * 8 + sg;
    if (t >= n_tok) return;
    device const float* l = logits + ulong(t) * ulong(n_exp);
    float v[16];
    for (int i = 0; i < 16; ++i) {
        const int idx = lane + 32 * i;
        v[i] = idx < n_exp ? l[idx] : -INFINITY;
    }
    float mx = -INFINITY;
    for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
    for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
    float sum = 0.0f;
    for (int i = 0; i < 16; ++i) {
        v[i] = (lane + 32 * i < n_exp) ? exp(v[i] - mx) : 0.0f;
        sum += v[i];
    }
    for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
    const float rcp = 1.0f / sum;
    for (int i = 0; i < 16; ++i) {
        if (lane + 32 * i >= n_exp) {
            v[i] = -INFINITY;
        } else {
            v[i] *= rcp;
            if (isnan(v[i])) v[i] = -FLT_MAX;
        }
    }"""
    assert old_ptopk in p_txt, "route_topk target not found in prompt.metal"
    p_path.write_text(p_txt.replace(old_ptopk, new_ptopk, 1), encoding="utf-8")


def patch_stream() -> None:
    sh_path = ROOT / "metal" / "src" / "engine" / "stream.h"
    sc_path = ROOT / "metal" / "src" / "engine" / "stream.cpp"

    sh_txt = sh_path.read_text(encoding="utf-8")
    sh_txt = sh_txt.replace(
        "uint8_t* region(uint64_t bytes);\n    void hold_region(bool hold);",
        "uint8_t* region(uint64_t bytes);\n    void hold_region(bool hold);\n    uint64_t release_region();\n    size_t region_stolen_slots() const { return region_slots_.size(); }",
    )
    sh_txt = sh_txt.replace(
        "struct Slot { int16_t layer = -1; int16_t e = -1; int prev = -1, next = -1; bool pre = false; int left = 0; uint64_t pin = 0; std::atomic<int> state{0}; };",
        "struct Slot { int16_t layer = -1; int16_t e = -1; int prev = -1, next = -1; bool pre = false, hot = false; int left = 0; uint64_t pin = 0; std::atomic<int> state{0}; };",
    )
    sh_txt = sh_txt.replace(
        "int nio_ = 6, pre_k_ = 10, pre_k2_ = 0, split_ = 2;",
        "int ne_ = 512, nio_ = 6, pre_k_ = 10, pre_k2_ = 0, split_ = 2;\n    float moce_bonus_ = 0.0f;\n    bool in_prompt_ = false;\n    uint64_t sep_region_bytes_ = 0;",
    )
    sh_path.write_text(sh_txt, encoding="utf-8")

    sc_txt = sc_path.read_text(encoding="utf-8")
    # 1. ExpertStream::open: dynamic ne_ from blk.0.ffn_gate_exps.weight, and optional spot check when dense-only gguf is used
    old_open = """    lay_.resize(NL);
    for (int l = 0; l < NL; ++l) {
        const std::string b = "blk." + std::to_string(l) + ".";
        const gguf::Tensor* g = f.find(b + "ffn_gate_exps.weight");
        const gguf::Tensor* u = f.find(b + "ffn_up_exps.weight");
        const gguf::Tensor* d = f.find(b + "ffn_down_exps.weight");
        if (!g || !u || !d) { err = "metal stream: layer " + std::to_string(l) + " has no routed experts"; return false; }
        lay_[(size_t) l] = {g->offset, g->bytes / NE, u->offset, u->bytes / NE, d->offset, d->bytes / NE, -1};
    }"""
    new_open = """    const gguf::Tensor* g0 = f.find("blk.0.ffn_gate_exps.weight");
    if (g0 && g0->shape.size() >= 3 && g0->shape[2] >= 1 && g0->shape[2] <= NE) {
        ne_ = (int) g0->shape[2];
    }
    if (const char* mb = std::getenv("STRATA_METAL_MOCE_BONUS")) moce_bonus_ = (float) std::atof(mb);
    lay_.resize(NL);
    for (int l = 0; l < NL; ++l) {
        const std::string b = "blk." + std::to_string(l) + ".";
        const gguf::Tensor* g = f.find(b + "ffn_gate_exps.weight");
        const gguf::Tensor* u = f.find(b + "ffn_up_exps.weight");
        const gguf::Tensor* d = f.find(b + "ffn_down_exps.weight");
        if (!g || !u || !d) { err = "metal stream: layer " + std::to_string(l) + " has no routed experts"; return false; }
        lay_[(size_t) l] = {g->offset, g->bytes / (uint64_t) ne_, u->offset, u->bytes / (uint64_t) ne_, d->offset, d->bytes / (uint64_t) ne_, -1};
    }"""
    assert old_open in sc_txt, "open target not found in stream.cpp"
    sc_txt = sc_txt.replace(old_open, new_open, 1)

    old_strx_chk = """            ok = ok && ver == 1 && nl == (uint32_t) NL && ne == (uint32_t) NE && src == (uint64_t) st.st_size;
            for (int l = 0; ok && l < NL; ++l) {
                const Lay& x = lay_[(size_t) l];
                ok = t[(size_t) l * 4 + 1] == x.gs && t[(size_t) l * 4 + 2] == x.us && t[(size_t) l * 4 + 3] == x.ds;
                off[(size_t) l] = t[(size_t) l * 4];
            }
            if (ok) {   // spot check: layer 0's first and layer 47's last expert, every slice
                for (const int l : {0, NL - 1}) {
                    const Lay& x = lay_[(size_t) l];
                    const int e = l == 0 ? 0 : NE - 1;"""
    new_strx_chk = """            ok = ok && ver == 1 && nl == (uint32_t) NL && ne == (uint32_t) ne_ && src == (uint64_t) st.st_size;
            for (int l = 0; ok && l < NL; ++l) {
                const Lay& x = lay_[(size_t) l];
                ok = t[(size_t) l * 4 + 1] == x.gs && t[(size_t) l * 4 + 2] == x.us && t[(size_t) l * 4 + 3] == x.ds;
                off[(size_t) l] = t[(size_t) l * 4];
            }
            const Lay& last_lay = lay_[(size_t) (NL - 1)];
            const bool has_gguf_exps = (uint64_t) st.st_size >= last_lay.d + (uint64_t) ne_ * last_lay.ds;
            if (ok && has_gguf_exps) {   // spot check: layer 0's first and layer 47's last expert, every slice
                for (const int l : {0, NL - 1}) {
                    const Lay& x = lay_[(size_t) l];
                    const int e = l == 0 ? 0 : ne_ - 1;"""
    assert old_strx_chk in sc_txt, "strx check not found in stream.cpp"
    sc_txt = sc_txt.replace(old_strx_chk, new_strx_chk, 1)

    # 2. Zero-copy file-backed load_dense with computed byte totals (no hard-coded sizes)
    old_load_dense = """bool ExpertStream::load_dense(gguf::File& f, mtl::Context& cx, const std::vector<std::string>& skip, mtl::Buffer& out,
                              std::string& err) {
    const double t0 = now_ms();
    std::vector<gguf::Tensor*> ts;
    uint64_t tot = 0;
    for (auto& kv : f.tensors) {
        gguf::Tensor& t = kv.second;
        if (t.name.find("_exps") != std::string::npos) continue;
        if (std::find(skip.begin(), skip.end(), t.name) != skip.end()) { t.offset = 0; continue; }
        ts.push_back(&t);
        tot += (t.bytes + 255) / 256 * 256;
    }"""
    new_load_dense = """bool ExpertStream::load_dense(gguf::File& f, mtl::Context& cx, const std::vector<std::string>& skip, mtl::Buffer& out,
                              std::string& err) {
    const double t0 = now_ms();
    uint64_t max_dense_end = 0;
    uint64_t min_exp_off = UINT64_MAX;
    size_t dense_cnt = 0;
    for (const auto& kv : f.tensors) {
        const gguf::Tensor& t = kv.second;
        if (t.name.find("_exps") != std::string::npos) {
            min_exp_off = std::min(min_exp_off, t.offset);
            continue;
        }
        ++dense_cnt;
        max_dense_end = std::max(max_dense_end, (t.offset + t.bytes + PAGE - 1) / PAGE * PAGE);
    }
    if (f.map != nullptr && max_dense_end > 0 && max_dense_end <= f.map_bytes && max_dense_end <= min_exp_off &&
        env_int("STRATA_METAL_MMAP_DENSE", 1) != 0) {
        const size_t alloc_before = cx.allocated();
        uint64_t locked_bytes = 0, skipped_hc_bytes = 0, unpinned_embd_bytes = 0;
        for (const auto& kv : f.tensors) {
            const gguf::Tensor& t = kv.second;
            if (t.name.find("_exps") != std::string::npos) continue;
            const uintptr_t p0 = ((uintptr_t) f.map + t.offset) / PAGE * PAGE;
            const uintptr_t p1 = ((uintptr_t) f.map + t.offset + t.bytes + PAGE - 1) / PAGE * PAGE;
            const bool is_skip = std::find(skip.begin(), skip.end(), t.name) != skip.end();
            const bool is_embd = t.name == "token_embd.weight";
            if (is_skip || is_embd) {
                if (is_skip) skipped_hc_bytes += t.bytes;
                if (is_embd) unpinned_embd_bytes += t.bytes;
                madvise((void*) p0, p1 - p0, MADV_DONTNEED);
            } else if (env_int("STRATA_METAL_MLOCK_DENSE", 1) != 0) {
                madvise((void*) p0, p1 - p0, MADV_WILLNEED);
                if (mlock((const void*) p0, p1 - p0) == 0) locked_bytes += (p1 - p0);
            }
        }
        if (f.map_bytes > max_dense_end) {
            madvise((uint8_t*) f.map + max_dense_end, f.map_bytes - (size_t) max_dense_end, MADV_DONTNEED);
        }
        out = cx.wrap(f.map, (size_t) max_dense_end, /*count_allocated=*/false);
        (void) alloc_before;
        std::fprintf(stderr, "strata metal stream: locked %.2f GB of active GPU dense tensors (unpinned %.2f GB CPU token_embd + %.2f GB skipped BF16 HC in %.2f GB file map)\\n",
                     locked_bytes / 1e9, unpinned_embd_bytes / 1e9, skipped_hc_bytes / 1e9, max_dense_end / 1e9);
        std::fprintf(stderr, "strata metal stream: %zu dense tensors wrapped zero-copy from file (%.2f GB clean file-backed map, %.1f ms)\\n",
                     dense_cnt, max_dense_end / 1e9, now_ms() - t0);
        return true;
    }
    std::vector<gguf::Tensor*> ts;
    uint64_t tot = 0;
    for (auto& kv : f.tensors) {
        gguf::Tensor& t = kv.second;
        if (t.name.find("_exps") != std::string::npos) continue;
        if (std::find(skip.begin(), skip.end(), t.name) != skip.end()) { t.offset = 0; continue; }
        ts.push_back(&t);
        tot += (t.bytes + 255) / 256 * 256;
    }"""
    assert old_load_dense in sc_txt, "load_dense target not found in stream.cpp"
    sc_txt = sc_txt.replace(old_load_dense, new_load_dense, 1)

    # 3. STRP profile loading & rank fallback using ne_ (while keeping where_ / smap_ indexed by l * NE + e where NE = 512)
    old_prof = """        if (std::fread(magic, 1, 4, f) == 4 && std::memcmp(magic, "STRP", 4) == 0 && std::fread(h, 4, 5, f) == 5 &&
            h[1] == (uint32_t) NL && h[2] == (uint32_t) NE) {
            for (uint32_t i = 0; i < h[4]; ++i) {
                uint16_t p[2];
                if (std::fread(p, 2, 2, f) != 2) break;
                if (p[0] < NL && p[1] < NE && !seen[(size_t) p[0] * NE + p[1]]) {
                    seen[(size_t) p[0] * NE + p[1]] = 1;
                    rank.push_back({p[0], p[1]});
                }
            }
        }
        std::fclose(f);
    } else {
        std::fprintf(stderr, "strata metal stream: no expert profile at %s: the arena starts layer-interleaved\\n", profile.c_str());
    }
    const size_t from_profile = rank.size();
    for (int e = 0; e < NE; ++e)
        for (int l = 0; l < NL; ++l)
            if (!seen[(size_t) l * NE + e]) rank.push_back({l, e});"""
    new_prof = """        if (std::fread(magic, 1, 4, f) == 4 && std::memcmp(magic, "STRP", 4) == 0 && std::fread(h, 4, 5, f) == 5 &&
            h[1] == (uint32_t) NL && h[2] == (uint32_t) ne_) {
            for (uint32_t i = 0; i < h[4]; ++i) {
                uint16_t p[2];
                if (std::fread(p, 2, 2, f) != 2) break;
                if (p[0] < NL && p[1] < (uint16_t) ne_ && !seen[(size_t) p[0] * NE + p[1]]) {
                    seen[(size_t) p[0] * NE + p[1]] = 1;
                    rank.push_back({p[0], p[1]});
                }
            }
        }
        std::fclose(f);
    } else {
        std::fprintf(stderr, "strata metal stream: no expert profile at %s: the arena starts layer-interleaved\\n", profile.c_str());
    }
    const size_t from_profile = rank.size();
    for (int e = 0; e < ne_; ++e)
        for (int l = 0; l < NL; ++l)
            if (!seen[(size_t) l * NE + e]) rank.push_back({l, e});"""
    assert old_prof in sc_txt, "prof target not found in stream.cpp"
    sc_txt = sc_txt.replace(old_prof, new_prof, 1)

    # 4. ahead(): top-3 unperturbed protection + smap residency consistency with GPU moe_route
    old_ahead = """void ExpertStream::ahead(int ln, int T, const float* pred, int K) {
    K = std::min(K, 32);
    std::pair<float, int> pw[8 * 32];
    int np = 0;
    for (int t = 0; t < T; ++t) {
        const float* lg = pred + (size_t) t * NE;
        int top[32];
        int nt = 0;
        for (int e = 0; e < NE; ++e) {
            if (nt < K) { top[nt++] = e; continue; }
            int m = 0;
            for (int j = 1; j < K; ++j) if (lg[top[j]] < lg[top[m]]) m = j;
            if (lg[e] > lg[top[m]]) top[m] = e;
        }
        for (int j = 0; j < nt; ++j) {
            int q = 0;
            while (q < np && pw[q].second != top[j]) ++q;
            if (q == np) pw[np++] = {lg[top[j]], top[j]};
            else pw[q].first = std::max(pw[q].first, lg[top[j]]);
        }
    }"""
    new_ahead = """void ExpertStream::ahead(int ln, int T, const float* pred, int K) {
    K = std::min(K, 32);
    static const bool unpert_all = env_int("STRATA_METAL_PREFETCH_UNPERT", 0) != 0;
    const uint32_t* sm = smap_.as<uint32_t>() + (size_t) ln * NE;
    std::pair<float, int> pw[8 * 32];
    int np = 0;
    for (int t = 0; t < T; ++t) {
        const float* lg = pred + (size_t) t * NE;
        auto eff_lg = [&](int e) -> float {
            return lg[e] + (moce_bonus_ > 0.0f && sm[e] != 0xFFFFFFFFu ? moce_bonus_ : 0.0f);
        };
        const int k_unpert = unpert_all ? K : std::min(3, K);
        int top[32];
        int nt = 0;
        for (int r = 0; r < k_unpert; ++r) {
            int best_e = -1;
            float best_v = -INFINITY;
            for (int e = 0; e < ne_; ++e) {
                bool seen_e = false;
                for (int j = 0; j < nt; ++j) if (top[j] == e) { seen_e = true; break; }
                if (!seen_e && lg[e] > best_v) { best_v = lg[e]; best_e = e; }
            }
            if (best_e >= 0) top[nt++] = best_e;
        }
        for (int e = 0; e < ne_; ++e) {
            if (k_unpert >= K) break;
            bool in_top = false;
            for (int j = 0; j < k_unpert; ++j) if (top[j] == e) { in_top = true; break; }
            if (in_top) continue;
            if (nt < K) { top[nt++] = e; continue; }
            int m = k_unpert;
            for (int j = k_unpert + 1; j < K; ++j) if (eff_lg(top[j]) < eff_lg(top[m])) m = j;
            if (eff_lg(e) > eff_lg(top[m])) top[m] = e;
        }
        for (int j = 0; j < nt; ++j) {
            const float prio = (j < k_unpert) ? (lg[top[j]] + moce_bonus_ + 1e3f) : eff_lg(top[j]);
            int q = 0;
            while (q < np && pw[q].second != top[j]) ++q;
            if (q == np) pw[np++] = {prio, top[j]};
            else pw[q].first = std::max(pw[q].first, prio);
        }
    }"""
    assert old_ahead in sc_txt, "ahead target not found in stream.cpp"
    sc_txt = sc_txt.replace(old_ahead, new_ahead, 1)

    # 5. victim() & multi-turn hot-slot protection + separate prompt scratch munmap in release_region()
    old_victim = """int ExpertStream::victim(int ci) {
    Cls& k = cls_[(size_t) ci];
    for (int s = k.tail; s >= 0; s = slots_[(size_t) s].prev)
        if (slots_[(size_t) s].state.load() == 0 && slots_[(size_t) s].pin != stamp_) {"""
    new_victim = """int ExpertStream::victim(int ci) {
    Cls& k = cls_[(size_t) ci];
    const int n_passes = in_prompt_ ? 2 : 1;
    for (int pass = 0; pass < n_passes; ++pass) {
        for (int s = k.tail; s >= 0; s = slots_[(size_t) s].prev)
            if (slots_[(size_t) s].state.load() == 0 && slots_[(size_t) s].pin != stamp_ && (!in_prompt_ || pass == 1 || !slots_[(size_t) s].hot)) {"""
    assert old_victim in sc_txt, "victim target not found in stream.cpp"
    sc_txt = sc_txt.replace(old_victim, new_victim, 1)
    sc_txt = sc_txt.replace(
        "            x.layer = x.e = -1;\n            return s;\n        }\n    return -1;",
        "            x.layer = x.e = -1;\n            return s;\n        }\n    }\n    return -1;",
        1,
    )

    old_lru_order = """        // the LRU order = the rank order (best at the front)
        for (size_t i = done; i-- > 0;) lru_front(where_[(size_t) (fill[i].first * NE + fill[i].second)]);"""
    new_lru_order = """        // the LRU order = the rank order (best at the front); mark top (k.n - (ne_ + 32)) STRP slots hot
        std::vector<int> cls_seen(cls_.size(), 0);
        for (size_t i = 0; i < done; ++i) {
            const int s = where_[(size_t) (fill[i].first * NE + fill[i].second)];
            if (s >= 0) {
                const int ci = cls_of(s);
                const int hot_cap = std::max(0, cls_[(size_t) ci].n - (ne_ + 32));
                slots_[(size_t) s].hot = (cls_seen[(size_t) ci] < hot_cap);
                ++cls_seen[(size_t) ci];
            }
        }
        for (size_t i = done; i-- > 0;) lru_front(where_[(size_t) (fill[i].first * NE + fill[i].second)]);"""
    assert old_lru_order in sc_txt, "lru_order target not found in stream.cpp"
    sc_txt = sc_txt.replace(old_lru_order, new_lru_order, 1)

    old_region_fn = """uint8_t* ExpertStream::region(uint64_t bytes) {"""
    new_region_fn = """uint8_t* ExpertStream::region(uint64_t bytes) {
    if (env_int("STRATA_METAL_SEPARATE_PROMPT_BUF", 1) != 0) {
        const uint64_t sz = (bytes + PAGE - 1) / PAGE * PAGE;
        void* m = mmap(nullptr, (size_t) sz, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
        if (m == MAP_FAILED) return nullptr;
        region_slots_.clear();
        region_ = (uint8_t*) m;
        region_bytes_ = 0;  // fresh zero pages from MAP_ANON; no arena slots to memset
        sep_region_bytes_ = sz;
        std::fprintf(stderr, "strata metal stream: prompt scratch %.2f GB mapped in separate anonymous mmap\\n",
                     sz / 1e9);
        return region_;
    }"""
    assert old_region_fn in sc_txt, "region_fn not found in stream.cpp"
    sc_txt = sc_txt.replace(old_region_fn, new_region_fn, 1)
    sc_txt = sc_txt.replace(
        "region_bytes_ = 0;   // fresh zero pages; nothing in the arena to empty",
        "region_bytes_ = 0; sep_region_bytes_ = sz;   // fresh zero pages; unmapped in release_region()",
    )

    old_hold_fn = """void ExpertStream::hold_region(bool hold) {
    std::unique_lock<std::mutex> g(mu_);
    if (hold) {
        // nothing may be reading into those slots: the queued read-aheads are dropped, the running ones waited for"""
    new_hold_fn = """void ExpertStream::hold_region(bool hold) {
    std::unique_lock<std::mutex> g(mu_);
    in_prompt_ = hold;
    if (hold) {
        // Refresh hot flags from current LRU order (protecting top k.n - (ne_ + 32) slots in every prompt turn)
        for (size_t ci = 0; ci < cls_.size(); ++ci) {
            const int hot_cap = std::max(0, cls_[ci].n - (ne_ + 32));
            int rank_idx = 0;
            for (int s = cls_[ci].head; s >= 0; s = slots_[(size_t) s].next) {
                slots_[(size_t) s].hot = (rank_idx < hot_cap);
                ++rank_idx;
            }
        }
        // nothing may be reading into those slots: the queued read-aheads are dropped, the running ones waited for"""
    assert old_hold_fn in sc_txt, "hold_fn start not found in stream.cpp"
    sc_txt = sc_txt.replace(old_hold_fn, new_hold_fn, 1)

    old_hold_end = """        if (region_) std::memset(region_, 0, (size_t) region_bytes_);
    } else {
        for (int s : region_slots_) {
            slots_[(size_t) s].state = 0;
            if (!retired_[(size_t) s]) lru_back(s);   // a retired slot's memory serves the prompt only
        }
    }
}"""
    new_hold_end = """        if (region_ && region_bytes_ > 0) std::memset(region_, 0, (size_t) region_bytes_);
    } else {
        for (int s : region_slots_) {
            slots_[(size_t) s].state = 0;
            if (!retired_[(size_t) s]) lru_back(s);   // a retired slot's memory serves the prompt only
        }
    }
}

uint64_t ExpertStream::release_region() {
    std::unique_lock<std::mutex> g(mu_);
    if (sep_region_bytes_ > 0 && region_ != nullptr) {
        const uint64_t freed = sep_region_bytes_;
        munmap(region_, (size_t) sep_region_bytes_);
        region_ = nullptr;
        sep_region_bytes_ = 0;
        return freed;
    }
    return 0;
}"""
    assert old_hold_end in sc_txt, "hold_fn end not found in stream.cpp"
    sc_txt = sc_txt.replace(old_hold_end, new_hold_end, 1)

    # Also during prompt chunks in at_layer, move streamed prompt slots to lru_back so victim() cycles the unpinned tail slots!
    sc_txt = sc_txt.replace(
        "submit(s, ln, e, false);\n                slots_[(size_t) s].pin = stamp_;\n                ++p_streamed_;",
        "submit(s, ln, e, false);\n                slots_[(size_t) s].pin = stamp_;\n                lru_back(s);\n                ++p_streamed_;",
    )

    sc_txt = sc_txt.replace("if (e < 0 || e >= NE || used[e]) continue;", "if (e < 0 || e >= ne_ || used[e]) continue;")
    sc_txt = sc_txt.replace("for (int e = 0; e < NE; ++e) {\n                if (where_[(size_t) (ln * NE + e)] >= 0) continue;",
                            "for (int e = 0; e < ne_; ++e) {\n                if (where_[(size_t) (ln * NE + e)] >= 0) continue;")
    sc_txt = sc_txt.replace("tight = tight || (cls_[i].n > 0 && cls_[i].n - stolen[i] < NE);",
                            "tight = tight || (cls_[i].n > 0 && cls_[i].n - stolen[i] < ne_);")

    # 6. Eliminate CPU/GPU data race on smap_ during background read-ahead:
    #    Background read-ahead threads (x.pre == true) never write to smap_ while the GPU is running moe_route.
    #    Instead, smap_ is written only when !x.pre (during make_arena or urgent/promoted waits while the GPU is paused)
    #    or inside at_layer() when a prefetched slot is claimed by the layer!
    old_io_smap = """            if (--x.left == 0) {
                smap_.as<uint32_t>()[j.layer * NE + j.e] = code_[(size_t) j.slot];
                x.state = 0;"""
    new_io_smap = """            if (--x.left == 0) {
                if (!x.pre) smap_.as<uint32_t>()[j.layer * NE + j.e] = code_[(size_t) j.slot];
                x.state = 0;"""
    assert old_io_smap in sc_txt, "old_io_smap not found in stream.cpp"
    sc_txt = sc_txt.replace(old_io_smap, new_io_smap, 1)

    old_claim = """                if (!(prompt && x.pre)) lru_front(s);   // a prompt's streamed experts stay at the LRU tail
                x.pre = false;"""
    new_claim = """                if (!(prompt && x.pre)) lru_front(s);   // a prompt's streamed experts stay at the LRU tail
                x.pre = false;
                if (st == 0) smap_.as<uint32_t>()[key] = code_[(size_t) s];"""
    assert old_claim in sc_txt, "old_claim not found in stream.cpp"
    sc_txt = sc_txt.replace(old_claim, new_claim, 1)

    old_drop_block = """            for (int s : drop) cancel_slot(s);
        }"""
    new_drop_block = """            for (int s : drop) cancel_slot(s);
        }
        if (moce_bonus_ > 0.0f) {
            std::unique_lock<std::mutex> g2(mu_, std::adopt_lock);
            done_cv_.wait(g2, [&] {
                for (int e = 0; e < ne_; ++e) {
                    if (used[e]) continue;
                    const int s = where_[(size_t) (l * NE + e)];
                    if (s >= 0 && slots_[(size_t) s].pre && slots_[(size_t) s].state.load() != 0) return false;
                }
                return true;
            });
            g2.release();
            // Evict any unchosen speculative read-ahead of this layer that finished so where_, smap_, and LRU order stay 100% deterministic
            for (int e = 0; e < ne_; ++e) {
                if (used[e]) continue;
                const int s = where_[(size_t) (l * NE + e)];
                if (s >= 0 && slots_[(size_t) s].pre && slots_[(size_t) s].state.load() == 0) {
                    Slot& x = slots_[(size_t) s];
                    where_[(size_t) (l * NE + e)] = -1;
                    smap_.as<uint32_t>()[l * NE + e] = 0xFFFFFFFFu;
                    x.layer = x.e = -1;
                    x.pre = false;
                    lru_back(s);
                }
            }
        }"""
    assert old_drop_block in sc_txt, "old_drop_block not found in stream.cpp"
    sc_txt = sc_txt.replace(old_drop_block, new_drop_block, 1)

    sc_path.write_text(sc_txt, encoding="utf-8")


def patch_window() -> None:
    w_path = ROOT / "metal" / "src" / "engine" / "window.cpp"
    pc_path = ROOT / "metal" / "src" / "engine" / "prompt_chain.inc"

    w_txt = w_path.read_text(encoding="utf-8")
    # 1. Add n_exp_, moce_pad_, cached pg shader lib IDs, and release_prompt_gemm declaration to Window
    old_seg_t = "    int seg_T_ = 0;"
    new_seg_t = "    int seg_T_ = 0, n_exp_ = NEXP;\n    uint32_t moce_pad_ = 0;\n    int pg_lib_ = -1, pg_lgdn_ = -1, pg_lqsa_ = -1;"
    assert old_seg_t in w_txt, "seg_T_ not found in window.cpp"
    w_txt = w_txt.replace(old_seg_t, new_seg_t, 1)

    w_txt = w_txt.replace(
        "void prompt_region(bool hold) { if (stream_ && st_ && st_->ready()) st_->hold_region(hold); }",
        "void prompt_region(bool hold);\n    uint64_t prompt_gemm_need_bytes() const;",
    )

    # 2. Allow any 1 <= g.n_expert <= NEXP in Window::init, and purge unused device allocations upfront
    old_init_chk = """    if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_ff != NFF || g.n_expert != NEXP || g.n_layers != 48) {
        err = "metal window: the geometry differs from Qwen3.8-Flash-Next's";
        return false;
    }"""
    new_init_chk = """    if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_ff != NFF || g.n_expert < 1 || g.n_expert > NEXP || g.n_layers != 48) {
        err = "metal window: the geometry differs from Qwen3.8-Flash-Next's";
        return false;
    }
    n_exp_ = (int) g.n_expert;
    if (const char* mb = std::getenv("STRATA_METAL_MOCE_BONUS")) {
        const float bonus = (float) std::atof(mb);
        if (bonus > 0.0f) {
            const float sc = std::exp(bonus);
            std::memcpy(&moce_pad_, &sc, 4);
        }
    }
    if (!(std::getenv("STRATA_METAL_KEEP_HOST_WEIGHTS") && std::string(std::getenv("STRATA_METAL_KEEP_HOST_WEIGHTS")) == "1")) {
        strata::gpu::purge_device_allocations(1u << 20);
    }"""
    assert old_init_chk in w_txt, "init check not found in window.cpp"
    w_txt = w_txt.replace(old_init_chk, new_init_chk, 1)

    # 3. Also detect n_exp_ from blk.0.ffn_gate_exps.weight in f_ right after f_.open
    old_fopen = """    if (!f_.open(shard1, err)) return false;"""
    new_fopen = """    if (!f_.open(shard1, err)) return false;
    if (const gguf::Tensor* ge0 = f_.find("blk.0.ffn_gate_exps.weight")) {
        if (ge0->shape.size() >= 3 && ge0->shape[2] >= 1 && ge0->shape[2] <= NEXP) n_exp_ = (int) ge0->shape[2];
    }"""
    assert old_fopen in w_txt, "f_.open not found in window.cpp"
    w_txt = w_txt.replace(old_fopen, new_fopen, 1)

    # 3b. Enable PLE when ss.ple.table is open (even without legacy CUDA host pointers)
    w_txt = w_txt.replace(
        'ple_on_ = ss.ple.ready() && f_.find("blk.1.ple_key.weight") != nullptr;',
        'ple_on_ = (ss.ple.ready() || (ss.ple.table != nullptr && ss.ple.table->is_open())) && f_.find("blk.1.ple_key.weight") != nullptr;',
    )
    # 3c. Add reset_stream_stats() helper to Window and window_reset_stats(void* w)
    w_txt = w_txt.replace(
        "void stream_report(FILE* f) const {",
        "void reset_stream_stats() { if (st_) st_->reset_stats(); for (int t = 0; t <= MAXT; ++t) { win_n_[t] = 0; win_ms_[t] = win_gpu_[t] = 0; } }\n    void stream_report(FILE* f) const {",
    )
    w_txt = w_txt.replace(
        "void window_reset() {",
        "void window_reset_stats(void* w) { if (w) ((Window*) w)->reset_stream_stats(); }\nvoid window_reset() {",
    )
    # 4. In Window::layer: use n_exp_, pass moce_pad_ and smap to router_route
    old_layer_route = """        a.n_tok = ntok;
        a.topk = TOPK;
        a.n_expert = NEXP;
        const bool dq2 = td.type == ref::Q2_0;
        const int gl = 32, gr = 1, gs = 4, ds_ = 4;
        // the down tile: Q2_0 8,1; else 4,2 - but on the M4 (no M5 layout) Q5_1 4,1 and Q8_0 8,1 (4,2 spills there:
        // 24-36 GB/s, mac-kern 513fc615)
        const bool m4 = !m5_layout();
        const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4,
                  dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2;
        a.gx_gu = NFF / (gs * (32 / gl) * gr);
        a.gx_dn = N / (ds_ * (32 / dl) * dr);
        mtl::Dispatch r;
        if (rf) {
            // the router GEMV + moe_route (bitwise both), beside the shared expert's GEMVs
            r.p = P(lib_, "router_route_" + std::to_string(trt.type), mtl::Constants().i(0, ntok), "router_route_" + std::to_string(trt.type) + "#" + std::to_string(ntok));
            GemvArgs ga{(uint32_t) N, 512, (uint32_t) ntok, (uint32_t) ref::row_bytes(trt.type, N), (uint32_t) N, 512};
            r.args = {mtl::Arg::B(W_, trt.offset), mtl::Arg::B(XF_.b, 0), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4),
                      mtl::Arg::B(logits_), mtl::Arg::V(ga), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_),
                      mtl::Arg::V(a), mtl::Arg::B(rcounter_)};
            r.grid[0] = 512 / 16;
            r.tg[0] = 512;
            nb_ = true;
        } else {
            r.p = P(lib_, "moe_route", {}, "moe_route");
            r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a)};
            r.tg[0] = 512;"""
    new_layer_route = """        a.n_tok = ntok;
        a.topk = TOPK;
        a.n_expert = (uint32_t) n_exp_;
        a.pad = stream_ ? moce_pad_ : 0u;
        const bool dq2 = td.type == ref::Q2_0;
        const int gl = 32, gr = 1, gs = 4, ds_ = 4;
        // the down tile: Q2_0 8,1; else 4,2 - but on the M4 (no M5 layout) Q5_1 4,1 and Q8_0 8,1 (4,2 spills there:
        // 24-36 GB/s, mac-kern 513fc615)
        const bool m4 = !m5_layout();
        const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4,
                  dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2;
        a.gx_gu = NFF / (gs * (32 / gl) * gr);
        a.gx_dn = N / (ds_ * (32 / dl) * dr);
        mtl::Dispatch r;
        const int mp_fc = std::getenv("STRATA_METAL_MOCE_PROTECT") ? std::max(0, std::min(10, std::atoi(std::getenv("STRATA_METAL_MOCE_PROTECT")))) : 3;
        uint32_t mr1_bits = 0x3F800000u; // 1.0f
        if (const char* mr1_env = std::getenv("STRATA_METAL_MOCE_BONUS_R1")) {
            const float mr1_b = (float) std::atof(mr1_env);
            if (mr1_b > 0.0f) {
                const float mr1_s = std::exp(mr1_b);
                std::memcpy(&mr1_bits, &mr1_s, 4);
            }
        }
        mtl::Constants r_fc;
        std::string r_fc_tag;
        if (mp_fc != 3) { r_fc.i(24, mp_fc); r_fc_tag += "#mp" + std::to_string(mp_fc); }
        if (mr1_bits != 0x3F800000u) { r_fc.u(25, mr1_bits); r_fc_tag += "#mr1_" + std::to_string(mr1_bits); }
        if (rf) {
            // the router GEMV + moe_route (bitwise both), beside the shared expert's GEMVs
            mtl::Constants rfc2 = r_fc;
            rfc2.i(0, ntok);
            r.p = P(lib_, "router_route_" + std::to_string(trt.type), rfc2, "router_route_" + std::to_string(trt.type) + "#" + std::to_string(ntok) + r_fc_tag);
            GemvArgs ga{(uint32_t) N, (uint32_t) n_exp_, (uint32_t) ntok, (uint32_t) ref::row_bytes(trt.type, N), (uint32_t) N, 512};
            r.args = {mtl::Arg::B(W_, trt.offset), mtl::Arg::B(XF_.b, 0), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4),
                      mtl::Arg::B(logits_), mtl::Arg::V(ga), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_),
                      mtl::Arg::V(a), mtl::Arg::B(rcounter_),
                      stream_ ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(rcounter_)};
            r.grid[0] = (uint32_t) ((n_exp_ + 15) / 16);
            r.tg[0] = 512;
            nb_ = true;
        } else {
            r.p = P(lib_, "moe_route", r_fc, "moe_route" + r_fc_tag);
            r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a),
                      stream_ ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(rcounter_)};
            r.tg[0] = 512;"""
    assert old_layer_route in w_txt, "layer_route not found in window.cpp"
    w_txt = w_txt.replace(old_layer_route, new_layer_route, 1)

    # 5. In stream_arena: min_slots = pg_on ? (n_exp_ + 64) : 64, account for transient PromptGemm scratch in plan log,
    #    and only pre-init PromptGemm inside stream_arena when NOT using a separate on-demand prompt buffer
    old_arena_tail = """    {   // the planned total: arena + the Metal buffers (or the process's own count when it is more) + the n-gram table's
        // reserve, against the cap (arena_budget.h: footprint_cap)
        const FootprintPlan fp = footprint_plan(am, budget);
        std::fprintf(stderr, "strata metal stream: planned resident footprint %.1f GB = arena %.1f + %.1f held before it + %.1f "
                             "reserved for the n-gram table; cap %.1f GB = the lesser of the GPU working set %.1f GB and the RAM "
                             "%.1f GB less %.1f GB left to macOS and its file cache%s\\n",
                     fp.total / gb, fp.arena / gb, fp.base / gb, fp.ple / gb, fp.cap / gb, cx_->recommended_working_set() / gb,
                     phys / gb, (am.phys - fp.cap) / gb,
                     am.arena_gb >= 0 && fp.over() ? " - STRATA_METAL_ARENA_GB is above the cap: kept as set (your choice), but "
                                                     "a Mac that runs short of memory can freeze (watchdog panics seen at 90+ GB "
                                                     "resident of 128)" : "");
    }
    const char* epg = std::getenv("STRATA_METAL_PROMPT_GEMM");
    const bool pg_on = !(epg && std::string(epg) == "0");
    bool minimal = false;   // the arena's least layout (the min slots) is not shrunk further: it goes ahead, filled in steps
    for (int tries = 0;; ++tries) {
        const double planned = budget;
        std::function<uint64_t()> recheck = [&]() -> uint64_t { return (uint64_t) arena_shrink(am, planned, avail_now()); };
        if (st_->make_arena(*cx_, (uint64_t) budget, prof, pg_on ? 576 : 64, err, minimal ? std::function<uint64_t()>() : recheck,
                            [&](uint64_t step) { return arena_grow(am, avail_now(), (double) step); }))
            break;
        if (st_->shrink_to() == 0 || tries == 4) return false;
        if ((double) st_->shrink_to() + 64.0 * (1 << 20) >= budget) minimal = true;   // no smaller layout than this one
        std::fprintf(stderr, "strata metal stream: less memory is free than when the arena was planned: %.1f -> %.1f GB\\n",
                     budget / gb, st_->shrink_to() / gb);
        budget = (double) st_->shrink_to();
    }
    if (pg_on && !prompt_gemm_init(err)) return false;"""
    new_arena_tail = """    const char* epg = std::getenv("STRATA_METAL_PROMPT_GEMM");
    const bool pg_on = !(epg && std::string(epg) == "0");
    const char* esep = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF");
    const bool sep_pg = !(esep && std::string(esep) == "0");
    const double pg_scratch_need = (pg_on && sep_pg) ? (double) prompt_gemm_need_bytes() : 0.0;
    if (am.arena_gb < 0 && pg_scratch_need > 0.0) {
        budget = std::max(1.0 * gb, budget - pg_scratch_need);
    }
    {   // the planned total: arena + the Metal buffers (or the process's own count when it is more) + the n-gram table's
        // reserve, against the cap (arena_budget.h: footprint_cap)
        const FootprintPlan fp = footprint_plan(am, budget);
        std::fprintf(stderr, "strata metal stream: planned resident footprint %.1f GB decode (%.1f GB virtual peak during prefill with %.2f GB transient prompt scratch) "
                             "= arena %.1f + %.1f held before it + %.1f reserved for the n-gram table; cap %.1f GB = the lesser of the GPU working set %.1f GB "
                             "and the RAM %.1f GB less %.1f GB left to macOS and its file cache%s\\n",
                     fp.total / gb, (fp.total + pg_scratch_need) / gb, pg_scratch_need / gb,
                     fp.arena / gb, fp.base / gb, fp.ple / gb, fp.cap / gb, cx_->recommended_working_set() / gb,
                     phys / gb, (am.phys - fp.cap) / gb,
                     am.arena_gb >= 0 && (fp.total + pg_scratch_need > fp.cap) ? " - peak prefill above cap" : "");
    }
    bool minimal = false;   // the arena's least layout (the min slots) is not shrunk further: it goes ahead, filled in steps
    for (int tries = 0;; ++tries) {
        const double planned = budget;
        std::function<uint64_t()> recheck = [&]() -> uint64_t { return (uint64_t) arena_shrink(am, planned, avail_now()); };
        if (st_->make_arena(*cx_, (uint64_t) budget, prof, pg_on ? (n_exp_ + 64) : 64, err, minimal ? std::function<uint64_t()>() : recheck,
                            [&](uint64_t step) { return arena_grow(am, avail_now(), (double) step); }))
            break;
        if (st_->shrink_to() == 0 || tries == 4) return false;
        if ((double) st_->shrink_to() + 64.0 * (1 << 20) >= budget) minimal = true;   // no smaller layout than this one
        std::fprintf(stderr, "strata metal stream: less memory is free than when the arena was planned: %.1f -> %.1f GB\\n",
                     budget / gb, st_->shrink_to() / gb);
        budget = (double) st_->shrink_to();
    }
    if (pg_on && !sep_pg && !prompt_gemm_init(err)) return false;"""
    assert old_arena_tail in w_txt, "arena_tail not found in window.cpp"
    w_txt = w_txt.replace(old_arena_tail, new_arena_tail, 1)

    # 6. In PromptGemm: cache compiled libraries across prompts, track all wrapped buffers in std::vector<mtl::Buffer*> bufs,
    #    and unmap + cx_->free them in Window::prompt_region(false) with measured proc_mem() footprint logging!
    old_pg_struct = """struct Window::PromptGemm {
    int MT = 0, QS = 256, TM = 16, maxrows = 0, nblk_max = 0;
    int lib = -1, lgdn = -1, lqsa = -1;
    bool k4 = false, gu_v3 = false, dn_v3 = false, flash = false;
    mtl::Buffer R, bo, inj[2], xnh, V, loh, mixed, mixh, qkv, z, al, be, ygh, gqn, gkn, gvc, ggb, gkq, go, yg, qfull, kcur, vcur,
        iqr, ikr, qcur, qcurh, qidx, scores, qids, part, XQ, xqh, logits, ids, wts, bcnt, boff, eoff, tiles, meta, stok, sw, pos,
        XS, HS, YE, yr, gsh, ush, hh, ys, sgv, emb, embh, pkey, pval, keyn, qn, gated, normW;
    mtl::Buffer WD, WU, Wqkv, Wz, Wal, Wbe, Wout, Wq, Wk, Wv, Wiq, Wik, Wr, Wsg, Wshg, Wshu, Wshd, Wpk, Wpv;
};"""
    new_pg_struct = """struct Window::PromptGemm {
    int MT = 0, QS = 256, TM = 16, maxrows = 0, nblk_max = 0;
    int lib = -1, lgdn = -1, lqsa = -1;
    bool k4 = false, gu_v3 = false, dn_v3 = false, flash = false;
    mtl::Buffer R, bo, inj[2], xnh, V, loh, mixed, mixh, qkv, z, al, be, ygh, gqn, gkn, gvc, ggb, gkq, go, yg, qfull, kcur, vcur,
        iqr, ikr, qcur, qcurh, qidx, scores, qids, part, XQ, xqh, logits, ids, wts, bcnt, boff, eoff, tiles, meta, stok, sw, pos,
        XS, HS, YE, yr, gsh, ush, hh, ys, sgv, emb, embh, pkey, pval, keyn, qn, gated, normW;
    mtl::Buffer WD, WU, Wqkv, Wz, Wal, Wbe, Wout, Wq, Wk, Wv, Wiq, Wik, Wr, Wsg, Wshg, Wshu, Wshd, Wpk, Wpv;
    std::vector<mtl::Buffer> all_bufs;
};"""
    assert old_pg_struct in w_txt, "pg_struct not found in window.cpp"
    w_txt = w_txt.replace(old_pg_struct, new_pg_struct, 1)

    w_txt = w_txt.replace("if (stream_ && !st_->ready()) return stream_arena(err);", "if (stream_ && !st_->ready() && !stream_arena(err)) return false;")
    old_pg_init_libs = """    p.lib = cx_->compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal", "gemv.metal",
                          "prompt.metal"});
    p.lgdn = cx_->compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\\n");
    p.lqsa = cx_->compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\\n");"""
    new_pg_init_libs = """    if (pg_lib_ < 0) {
        pg_lib_ = cx_->compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal", "gemv.metal",
                                "prompt.metal"});
        pg_lgdn_ = cx_->compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\\n");
        pg_lqsa_ = cx_->compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\\n");
    }
    p.lib = pg_lib_;
    p.lgdn = pg_lgdn_;
    p.lqsa = pg_lqsa_;"""
    assert old_pg_init_libs in w_txt, "pg_init_libs not found in window.cpp"
    w_txt = w_txt.replace(old_pg_init_libs, new_pg_init_libs, 1)

    old_ab = """        if (region == nullptr) return cx_->alloc(bytes);
        mtl::Buffer r = cx_->wrap(region, b);
        region += b;
        return r;"""
    new_ab = """        if (region == nullptr) {
            mtl::Buffer r = cx_->alloc(bytes);
            p.all_bufs.push_back(r);
            return r;
        }
        mtl::Buffer r = cx_->wrap(region, b, /*count_allocated=*/false);
        p.all_bufs.push_back(r);
        region += b;
        return r;"""
    assert old_ab in w_txt, "old_ab not found in window.cpp"
    w_txt = w_txt.replace(old_ab, new_ab, 1)

    old_pg_end = """    pg_ = pg.release();
    std::fprintf(stderr, "strata metal: prompts on the matrix units in chunks of %d (experts %s, QSA %s, matmul2d layout %s; "
                         "%.2f GB of buffers%s)\\n", C, p.gu_v3 ? "v3" : "v2", p.flash ? "flash" : "v1", p.k4 ? "K4" : "other",
                 need / 1e9, stream_ ? ", in the expert arena while a prompt runs" : "");
    (void) err;
    return true;
}"""
    new_pg_end = """    if (need != prompt_gemm_need_bytes()) { err = "metal window: prompt_gemm_need_bytes() != need"; return false; }
    pg_ = pg.release();
    const char* esep_log = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF");
    const bool sep_pg_log = !(esep_log && std::string(esep_log) == "0");
    std::fprintf(stderr, "strata metal: prompts on the matrix units in chunks of %d (experts %s, QSA %s, matmul2d layout %s; "
                         "%.2f GB of buffers%s)\\n", C, p.gu_v3 ? "v3" : "v2", p.flash ? "flash" : "v1", p.k4 ? "K4" : "other",
                 need / 1e9, stream_ ? (sep_pg_log ? ", in a separate transient mmap while a prompt runs" : ", in the expert arena while a prompt runs") : "");
    (void) err;
    return true;
}

uint64_t Window::prompt_gemm_need_bytes() const {
    const char* ec = std::getenv("STRATA_METAL_PROMPT_CHUNK");
    uint64_t phys = 0;
    size_t plen = sizeof phys;
    sysctlbyname("hw.memsize", &phys, &plen, nullptr, 0);
    const int stream_chunk = phys >= (96ull << 30) ? 4096 : 1024;
    const int C = std::max(64, ec ? std::atoi(ec) : (stream_ ? stream_chunk : 2048)) / 64 * 64;
    const int QS = std::min(C, 256);
    const bool k4 = ::coop_k4(*cx_);
    const char* eq = std::getenv("STRATA_METAL_PROMPT_QSA");
    const bool flash = eq ? std::string(eq) == "flash" : k4;
    PromptGemm p_tm;
    const size_t MT = (size_t) C;
    const size_t mr = (size_t) (C * TOPK + NEXP * p_tm.TM);
    const size_t nblk_max = (size_t) ((C * TOPK + 1023) / 1024);
    const size_t MB = (size_t) (max_cells_ / 4 + 2);
    uint64_t need = 0;
    auto AB = [&](size_t bytes) { need += (bytes + 16383) / 16384 * 16384; };
    auto F32 = [&](size_t n) { AB(4 * n); };
    auto F16 = [&](size_t n) { AB(2 * n); };
    F32(MT * HC * N); F32(MT * N); F32(MT * HC); F32(MT * HC); F16(MT * HC * N);
    F32(MT * (LR + HC)); F16(MT * LR); F32(MT * N); F16(MT * N);
    F32(MT * 10240); F32(MT * 6144); F32(MT * 48); F32(MT * 48); F32(MT * 6144);
    F16(MT * 6144); F32(MT * 16 * 128); F32(MT * 16 * 128); F32(MT * 6144);
    F32(MT * 96); F32(MT * 16); F32(MT * 6144);
    F32(MT * 12288); F32(MT * 512); F32(MT * 512); F32(MT * 512); F32(MT * 128);
    F32(MT * 24 * 256); F16(MT * 24 * 256); F32(MT * 4 * 128); F32(MT * MB);
    AB(4 * MT * 2051);
    if (flash) F32(16); else F32((size_t) QS * 33 * 24 * 258);
    F32(4 * MT * 6144 + 3 * (size_t) QS * 6144 + MT * 384); F16(MT * 6144);
    F32(MT * NEXP); AB(4 * MT * TOPK); F32(MT * TOPK);
    AB(4 * nblk_max * NEXP); AB(4 * nblk_max * NEXP);
    AB(4 * NEXP); AB(8 * (mr / p_tm.TM + 1)); AB(64);
    AB(4 * mr); F32(mr); AB(4 * MT * TOPK); F16(mr * N); F16(mr * NFF);
    F32(mr * N); F32(MT * N); F32(MT * NFF); F32(MT * NFF); F16(MT * NFF);
    F32(MT * N); F32(MT);
    F16((size_t) (LR + HC) * HC * N); F16((size_t) HC * N * LR); F16((size_t) 10240 * N);
    F16((size_t) 6144 * N); F16((size_t) 48 * N); F16((size_t) 48 * N); F16((size_t) N * 6144);
    F16((size_t) 12288 * N); F16((size_t) 512 * N); F16((size_t) 512 * N); F16((size_t) 512 * N);
    F16((size_t) 128 * N); F16((size_t) NEXP * N); F16((size_t) N); F16((size_t) NFF * N);
    F16((size_t) NFF * N); F16((size_t) N * NFF);
    if (ple_on_) {
        F32(MT * N); F16(MT * N); F32(MT * 10240); F32(MT * N); F32(MT * 10240);
        F32(MT * 10240); F32(MT * 10240); F32(MT * 10240);
        F16((size_t) 10240 * N); F16((size_t) N * N);
    }
    return need;
}

void Window::prompt_region(bool hold) {
    if (!stream_ || !st_ || !st_->ready()) return;
    if (hold) {
        st_->hold_region(true);
    } else {
        st_->hold_region(false);
        const char* esep = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF");
        const bool sep_pg = !(esep && std::string(esep) == "0");
        if (sep_pg && pg_ != nullptr) {
            const ProcMem pm_before = proc_mem();
            mtp_src_ = nullptr;
            for (mtl::Buffer& b : pg_->all_bufs) cx_->free(b);
            delete pg_;
            pg_ = nullptr;
            const uint64_t unmapped = st_->release_region();
            strata::gpu::purge_device_allocations(1u << 20);
            const ProcMem pm_after = proc_mem();
            if (unmapped > 0) {
                std::fprintf(stderr, "strata metal stream: prompt scratch %.2f GB unmapped after prefill (phys_footprint %.2f -> %.2f GB, resident RSS %.2f -> %.2f GB)\\n",
                             unmapped / 1e9, pm_before.footprint / 1e9, pm_after.footprint / 1e9, pm_before.rss / 1e9, pm_after.rss / 1e9);
            }
        }
    }
}"""
    assert old_pg_end in w_txt, "old_pg_end not found in window.cpp"
    w_txt = w_txt.replace(old_pg_end, new_pg_end, 1)

    old_pg_moe = """        gemm(p.mixh, p.Wr, p.logits, T, NEXP, N);
        const uint32_t nt = (uint32_t) T;
        D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits), mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::V(nt)}, (T + 7) / 8, 1, 256);
        const int nent = T * TOPK, nblk = (nent + 1023) / 1024;
        PgPlan pa{(uint32_t) nent, TOPK, NEXP, (uint32_t) p.TM, (uint32_t) nblk};
        const uint32_t nclr = (uint32_t) p.maxrows;
        D(PP(p.lib, "mm_clear"), {mtl::Arg::B(p.stok), mtl::Arg::V(nclr)}, (p.maxrows + 255) / 256, 1, 256);
        D(PP(p.lib, "mm_hist"), {mtl::Arg::B(p.ids), mtl::Arg::B(p.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024);
        D(PP(p.lib, "mm_scan"), {mtl::Arg::B(p.bcnt), mtl::Arg::B(p.boff), mtl::Arg::B(p.eoff), mtl::Arg::B(p.tiles),
                                 mtl::Arg::B(p.meta), mtl::Arg::V(pa)}, 1, 1, NEXP);"""
    new_pg_moe = """        gemm(p.mixh, p.Wr, p.logits, T, n_exp_, N);
        const uint32_t nt = (uint32_t) T | ((uint32_t) n_exp_ << 16);
        D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits), mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::V(nt)}, (T + 7) / 8, 1, 256);
        const int nent = T * TOPK, nblk = (nent + 1023) / 1024;
        PgPlan pa{(uint32_t) nent, TOPK, (uint32_t) n_exp_, (uint32_t) p.TM, (uint32_t) nblk};
        const uint32_t nclr = (uint32_t) p.maxrows;
        D(PP(p.lib, "mm_clear"), {mtl::Arg::B(p.stok), mtl::Arg::V(nclr)}, (p.maxrows + 255) / 256, 1, 256);
        D(PP(p.lib, "mm_hist"), {mtl::Arg::B(p.ids), mtl::Arg::B(p.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024);
        D(PP(p.lib, "mm_scan"), {mtl::Arg::B(p.bcnt), mtl::Arg::B(p.boff), mtl::Arg::B(p.eoff), mtl::Arg::B(p.tiles),
                                 mtl::Arg::B(p.meta), mtl::Arg::V(pa)}, 1, 1, (uint32_t) n_exp_);"""
    assert old_pg_moe in w_txt, "pg_moe target not found in window.cpp"
    w_txt = w_txt.replace(old_pg_moe, new_pg_moe, 1)
    w_txt = w_txt.replace(
        "r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a)};",
        "r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(rcounter_)};",
    )
    old_at_cut = """    void at_cut(int s) {
        st_->at_layer(s, seg_T_, TOPK, cut_ids_ ? cut_ids_ : ids_.as<int32_t>(), st_pred_ && !cut_ids_ ? predL_.as<float>() + (size_t) s * MAXT * 512 : nullptr,
                      predL2_.ptr && !cut_ids_ ? predL2_.as<float>() + (size_t) s * MAXT * 512 : nullptr);
    }"""
    new_at_cut = """    void dump_moe_layer(const char* pfx, int l, bool dump_in, bool dump_out);
    void at_cut(int s) {
        if (cut_ids_ && pg_) {
            if (const char* pfx = std::getenv("STRATA_METAL_DUMP_MOE")) {
                if (s > 0) dump_moe_layer(pfx, s - 1, false, true);
                dump_moe_layer(pfx, s, true, false);
            }
        }
        st_->at_layer(s, seg_T_, TOPK, cut_ids_ ? cut_ids_ : ids_.as<int32_t>(), st_pred_ && !cut_ids_ ? predL_.as<float>() + (size_t) s * MAXT * 512 : nullptr,
                      predL2_.ptr && !cut_ids_ ? predL2_.as<float>() + (size_t) s * MAXT * 512 : nullptr);
    }"""
    assert old_at_cut in w_txt, "old_at_cut not found in window.cpp"
    w_txt = w_txt.replace(old_at_cut, new_at_cut, 1)

    old_p_rows = "int Window::prompt_chunk_rows() const { return pg_ ? pg_->MT : 0; }"
    new_p_rows = """int Window::prompt_chunk_rows() const { return pg_ ? pg_->MT : 0; }
void Window::dump_moe_layer(const char* pfx, int l, bool dump_in, bool dump_out) {
    if (!pfx || !pg_) return;
    static std::string s_pfx;
    static size_t s_prev_T = 0;
    const size_t T = (size_t) seg_T_;
    const bool first_chunk = (last_T_ == 0 && s_pfx != pfx);
    if (l == 0 && dump_in && first_chunk) {
        s_pfx = pfx;
        s_prev_T = 0;
    }
    auto wr = [&](const char* ext, const void* ptr, size_t row_bytes) {
        std::string fn = std::string(pfx) + ext;
        const size_t chunk_bytes = T * row_bytes;
        if (s_prev_T == 0) {
            FILE* fp = std::fopen(fn.c_str(), l == 0 ? "wb" : "r+b");
            if (!fp) fp = std::fopen(fn.c_str(), "wb");
            if (!fp) { std::perror(("dump_moe fopen " + fn).c_str()); std::exit(2); }
            if (std::fseek(fp, (long)((size_t) l * chunk_bytes), SEEK_SET) != 0 ||
                std::fwrite(ptr, 1, chunk_bytes, fp) != chunk_bytes) {
                std::perror(("dump_moe fwrite " + fn).c_str()); std::exit(2);
            }
            std::fclose(fp);
        } else {
            // Multi-chunk prompt: splice new T rows into [NL, s_prev_T + T, row_bytes]
            const size_t old_l_bytes = s_prev_T * row_bytes;
            const size_t new_l_bytes = (s_prev_T + T) * row_bytes;
            static std::map<std::string, std::vector<uint8_t>> prev_files;
            std::vector<uint8_t>& prev_file = prev_files[ext];
            if (l == 0) {
                FILE* rf = std::fopen(fn.c_str(), "rb");
                if (!rf) { std::perror(("dump_moe multi-chunk r " + fn).c_str()); std::exit(2); }
                prev_file.resize((size_t) 48 * old_l_bytes);
                if (std::fread(prev_file.data(), 1, prev_file.size(), rf) != prev_file.size()) {
                    std::perror(("dump_moe multi-chunk fread " + fn).c_str()); std::exit(2);
                }
                std::fclose(rf);
            }
            FILE* wf = std::fopen(fn.c_str(), l == 0 ? "wb" : "r+b");
            if (!wf) { std::perror(("dump_moe multi-chunk w " + fn).c_str()); std::exit(2); }
            std::fseek(wf, (long)((size_t) l * new_l_bytes), SEEK_SET);
            if (std::fwrite(prev_file.data() + (size_t) l * old_l_bytes, 1, old_l_bytes, wf) != old_l_bytes ||
                std::fwrite(ptr, 1, chunk_bytes, wf) != chunk_bytes) {
                std::perror(("dump_moe multi-chunk fwrite " + fn).c_str()); std::exit(2);
            }
            std::fclose(wf);
            if (l == 47) prev_file.clear();
        }
    };
    if (dump_in) {
        wr(".mixed", pg_->mixed.ptr, N * sizeof(float));
        wr(".ids",   pg_->ids.ptr,   TOPK * sizeof(int32_t));
        wr(".wts",   pg_->wts.ptr,   TOPK * sizeof(float));
    }
    if (dump_out) {
        wr(".yr",  pg_->yr.ptr,  N * sizeof(float));
        wr(".hh",  pg_->hh.ptr,  NFF * sizeof(uint16_t));
        wr(".sgv", pg_->sgv.ptr, sizeof(float));
        if (l == 47) s_prev_T += T;
    }
}"""
    assert old_p_rows in w_txt, "old_p_rows not found in window.cpp"
    w_txt = w_txt.replace(old_p_rows, new_p_rows, 1)

    old_run_seg = """        seg_T_ = T;
        cut_ids_ = p.ids.as<int32_t>();
        cx_->run_graph_seg(ds, cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this);
        cut_ids_ = nullptr;"""
    new_run_seg = """        seg_T_ = T;
        cut_ids_ = p.ids.as<int32_t>();
        cx_->run_graph_seg(ds, cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this);
        if (const char* pfx = std::getenv("STRATA_METAL_DUMP_MOE")) {
            dump_moe_layer(pfx, 47, false, true);
        }
        cut_ids_ = nullptr;"""
    assert old_run_seg in w_txt, "old_run_seg not found in window.cpp"
    w_txt = w_txt.replace(old_run_seg, new_run_seg, 1)
    w_txt = w_txt.replace("if (kt == Tn && head && ++seen == 20 && !stream_) kprof(c.ds, Tn);", "if (kt == Tn && head && ++seen == 20) kprof(c.ds, Tn);")
    w_path.write_text(w_txt, encoding="utf-8")

    # 7. prompt_chain.inc: same packed nt and n_exp_ for gemm, route_topk, PcMmPlanArgs, and mm_scan
    pc_txt = pc_path.read_text(encoding="utf-8")
    old_pc_moe = """        gemm(b.mixh, b.Wr, b.logits, nT, NEXP, N);
        const uint32_t nt = (uint32_t) nT;
        D(P_(b.lib, "route_topk"), {mtl::Arg::B(b.logits), mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::V(nt)}, (nT + 7) / 8, 1, 256);
        const int nent = nT * TOPK, nblk = (nent + 1023) / 1024;
        PcMmPlanArgs pa{(uint32_t) nent, TOPK, NEXP, (uint32_t) PC_TM, (uint32_t) nblk};
        const uint32_t nclr = (uint32_t) b.maxrows;
        D(P_(b.lib, "mm_clear"), {mtl::Arg::B(b.stok), mtl::Arg::V(nclr)}, (b.maxrows + 255) / 256, 1, 256);
        D(P_(b.lib, "mm_hist"), {mtl::Arg::B(b.ids), mtl::Arg::B(b.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024);
        D(P_(b.lib, "mm_scan"), {mtl::Arg::B(b.bcnt), mtl::Arg::B(b.boff), mtl::Arg::B(b.eoff), mtl::Arg::B(b.tiles), mtl::Arg::B(b.meta),
                                mtl::Arg::V(pa)}, 1, 1, NEXP);"""
    new_pc_moe = """        gemm(b.mixh, b.Wr, b.logits, nT, n_exp_, N);
        const uint32_t nt = (uint32_t) nT | ((uint32_t) n_exp_ << 16);
        D(P_(b.lib, "route_topk"), {mtl::Arg::B(b.logits), mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::V(nt)}, (nT + 7) / 8, 1, 256);
        const int nent = nT * TOPK, nblk = (nent + 1023) / 1024;
        PcMmPlanArgs pa{(uint32_t) nent, TOPK, (uint32_t) n_exp_, (uint32_t) PC_TM, (uint32_t) nblk};
        const uint32_t nclr = (uint32_t) b.maxrows;
        D(P_(b.lib, "mm_clear"), {mtl::Arg::B(b.stok), mtl::Arg::V(nclr)}, (b.maxrows + 255) / 256, 1, 256);
        D(P_(b.lib, "mm_hist"), {mtl::Arg::B(b.ids), mtl::Arg::B(b.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024);
        D(P_(b.lib, "mm_scan"), {mtl::Arg::B(b.bcnt), mtl::Arg::B(b.boff), mtl::Arg::B(b.eoff), mtl::Arg::B(b.tiles), mtl::Arg::B(b.meta),
                                 mtl::Arg::V(pa)}, 1, 1, (uint32_t) n_exp_);"""
    assert old_pc_moe in pc_txt, "pc_moe target not found in prompt_chain.inc"
    pc_path.write_text(pc_txt.replace(old_pc_moe, new_pc_moe, 1), encoding="utf-8")


def patch_tests_and_route_parity() -> None:
    # 1. Any test in metal/tests that dispatches moe_route or router_route needs the extra smap buffer argument
    for p in (ROOT / "metal" / "tests").glob("*.cpp"):
        txt = p.read_text(encoding="utf-8")
        if "moe_route" in txt and "mtl::Arg::B(plan), mtl::Arg::V(a)};" in txt:
            txt = txt.replace(
                "mtl::Arg::B(plan), mtl::Arg::V(a)};",
                "mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(plan)};",
            )
            p.write_text(txt, encoding="utf-8")
        if "router_route_" in txt and "Arg::B(rcounter)};" in txt:
            txt = txt.replace(
                "mtl::Arg::V(a), mtl::Arg::B(rcounter)};",
                "mtl::Arg::V(a), mtl::Arg::B(rcounter), mtl::Arg::B(rcounter)};",
            )
            p.write_text(txt, encoding="utf-8")

    # 2. Extend route_parity.cpp to test BOTH K=512 (pad=0) AND K=256 + MoCE (pad=exp(0.85f), random smap, out-of-range logits in 256..511)
    rp_path = ROOT / "metal" / "tests" / "route_parity.cpp"
    rp_src = r"""// metal/tests/route_parity.cpp - moe_route (softmax + top-10 + the expert plan) against a CPU replica of
// native_router.cu's rule, for 1..8 tokens:
//   Phase 1: 200 windows with n_expert = 512, pad = 0 (deliberate exact ties)
//   Phase 2: 200 windows with n_expert = 256, pad = bit_cast(exp(0.85f)), random smap residency, and huge traps in 256..511
#include "mtl_rt.h"
#include "ref_quant.h"

#include <algorithm>
#include <cfloat>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <set>
#include <vector>

namespace {
constexpr int P_NG = 0, P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, P_ETOK = P_GSTART + P_MAXE + 4,
              P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE;
struct MoeArgs {
    uint64_t gate_off, up_off, down_off, gu_stride, d_stride;
    uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad;
};
}  // namespace

int main() {
    mtl::Context cx;
    const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal"});
    mtl::Pipeline p = cx.pipeline(lib, "moe_route");
    mtl::Buffer lg = cx.alloc(4 * 8 * 512), ids = cx.alloc(4 * 80), w = cx.alloc(4 * 80), plan = cx.alloc(4 * P_INTS),
                smap = cx.alloc(4 * 512);
    ref::Rng r(9);
    int fails_512 = 0, fails_256_moce = 0;
    double worst_w_512 = 0, worst_w_256_moce = 0;

    // Phase 1: K=512, MoCE off (pad = 0)
    for (int trial = 0; trial < 200; ++trial) {
        const int nt = 1 + trial % 8;
        for (int i = 0; i < nt * 512; ++i) {
            float v = r.sym() * 4.0f;
            if (trial % 3 == 0) v = std::round(v * 4.0f) / 4.0f;   // many exact ties
            lg.as<float>()[i] = v;
        }
        MoeArgs a{};
        a.n_tok = nt;
        a.topk = 10;
        a.n_expert = 512;
        a.gx_gu = 1;
        a.gx_dn = 1;
        a.pad = 0;
        mtl::Dispatch d;
        d.p = &p;
        d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(smap)};
        d.tg[0] = 512;
        cx.run({d});
        for (int t = 0; t < nt; ++t) {
            const float* l = lg.as<float>() + t * 512;
            float mx = -INFINITY;
            for (int e = 0; e < 512; ++e) mx = std::max(mx, l[e]);
            std::vector<double> pr(512);
            double s = 0;
            for (int e = 0; e < 512; ++e) s += (pr[e] = std::exp((double)l[e] - mx));
            for (int e = 0; e < 512; ++e) pr[e] /= s;
            std::vector<int> order(512);
            for (int e = 0; e < 512; ++e) order[e] = e;
            std::stable_sort(order.begin(), order.end(), [&](int x, int y) { return l[x] > l[y]; });
            double ssum = 0;
            for (int k = 0; k < 10; ++k) ssum += pr[order[k]];
            for (int k = 0; k < 10; ++k) {
                if (ids.as<int>()[t * 10 + k] != order[k]) { ++fails_512; break; }
                const double ref = pr[order[k]] / std::max(ssum, 6.103515625e-5);
                worst_w_512 = std::max(worst_w_512, std::fabs(w.as<float>()[t * 10 + k] - ref) / ref);
            }
        }
        const int* P = plan.as<int>();
        std::set<int> distinct(ids.as<int>(), ids.as<int>() + nt * 10);
        bool ok = P[P_NG] == (int)distinct.size();
        int gi = 0;
        for (int e : distinct) {
            if (P[P_GEXP + gi] != e) ok = false;
            int j = P[P_GSTART + gi];
            for (int i = 0; i < nt * 10; ++i)
                if (ids.as<int>()[i] == e) {
                    if (P[P_ETOK + j] != i / 10 || P[P_TOKENT + i] != j) ok = false;
                    ++j;
                }
            if (j != P[P_GSTART + gi + 1]) ok = false;
            ++gi;
        }
        fails_512 += !ok;
    }

    // Phase 2: K=256, MoCE on (pad = bit_cast(exp(0.85f))), random smap residency, and huge trap logits in 256..511
    const float moce_scale = std::exp(0.85f);
    uint32_t moce_pad = 0;
    std::memcpy(&moce_pad, &moce_scale, 4);
    for (int trial = 0; trial < 200; ++trial) {
        const int nt = 1 + trial % 8;
        for (int e = 0; e < 512; ++e) {
            smap.as<uint32_t>()[e] = (r.sym() > 0.0f) ? (uint32_t) e : 0xFFFFFFFFu;
        }
        for (int t = 0; t < nt; ++t) {
            for (int e = 0; e < 256; ++e) {
                float v = r.sym() * 4.0f;
                if (trial % 3 == 0) v = std::round(v * 4.0f) / 4.0f;
                lg.as<float>()[t * 512 + e] = v;
            }
            for (int e = 256; e < 512; ++e) {
                lg.as<float>()[t * 512 + e] = 50.0f + (float) e;  // trap logits past n_expert = 256
            }
        }
        MoeArgs a{};
        a.n_tok = nt;
        a.topk = 10;
        a.n_expert = 256;
        a.gx_gu = 1;
        a.gx_dn = 1;
        a.pad = moce_pad;
        mtl::Dispatch d;
        d.p = &p;
        d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(smap)};
        d.tg[0] = 512;
        cx.run({d});

        const uint32_t* sm = smap.as<uint32_t>();
        for (int t = 0; t < nt; ++t) {
            const float* l = lg.as<float>() + t * 512;
            float mx = -INFINITY;
            for (int e = 0; e < 256; ++e) mx = std::max(mx, l[e]);
            std::vector<float> pr(256);
            float s = 0.0f;
            for (int e = 0; e < 256; ++e) s += (pr[e] = std::exp(l[e] - mx));
            const float rcp = 1.0f / s;
            for (int e = 0; e < 256; ++e) pr[e] *= rcp;

            std::vector<float> work = pr;
            std::vector<int> picked(10, -1);
            double ssum = 0.0;
            for (int rank = 0; rank < 10; ++rank) {
                if (rank == 3) {
                    for (int e = 0; e < 256; ++e) {
                        if (work[e] > 0.0f && sm[e] != 0xFFFFFFFFu) work[e] *= moce_scale;
                    }
                }
                int best_e = 0;
                float best_v = work[0];
                for (int e = 1; e < 256; ++e) {
                    if (work[e] > best_v) { best_v = work[e]; best_e = e; }
                }
                picked[rank] = best_e;
                work[best_e] = -INFINITY;
                ssum += pr[best_e];
            }
            for (int k = 0; k < 10; ++k) {
                const int got_id = ids.as<int>()[t * 10 + k];
                if (got_id != picked[k] || got_id >= 256) { ++fails_256_moce; break; }
                const double ref = (double) pr[picked[k]] / std::max(ssum, 6.103515625e-5);
                worst_w_256_moce = std::max(worst_w_256_moce, std::fabs(w.as<float>()[t * 10 + k] - ref) / ref);
            }
        }
        const int* P = plan.as<int>();
        std::set<int> distinct(ids.as<int>(), ids.as<int>() + nt * 10);
        bool ok = P[P_NG] == (int)distinct.size();
        int gi = 0;
        for (int e : distinct) {
            if (e >= 256 || P[P_GEXP + gi] != e) ok = false;
            int j = P[P_GSTART + gi];
            for (int i = 0; i < nt * 10; ++i)
                if (ids.as<int>()[i] == e) {
                    if (P[P_ETOK + j] != i / 10 || P[P_TOKENT + i] != j) ok = false;
                    ++j;
                }
            if (j != P[P_GSTART + gi + 1]) ok = false;
            ++gi;
        }
        fails_256_moce += !ok;
    }

    const int total_fails = fails_512 + fails_256_moce;
    const double worst_w = std::max(worst_w_512, worst_w_256_moce);
    std::printf("route parity K=512 (pad=0, 200 windows): mismatches %d, worst weight rel err %.2e\n",
                fails_512, worst_w_512);
    std::printf("route parity K=256 + MoCE=0.85 (random smap + 256..511 traps, 200 windows): mismatches %d, worst weight rel err %.2e  %s\n",
                fails_256_moce, worst_w_256_moce, total_fails == 0 && worst_w < 2e-6 ? "PASS" : "FAIL");
    return total_fails == 0 && worst_w < 2e-6 ? 0 : 1;
}
"""
    rp_path.write_text(rp_src, encoding="utf-8")


def write_strata_window_cli_and_converter() -> None:
    cli_src = r"""// metal/src/strata_window_cli.cpp - Direct zero-bloat Metal runner for Strata's Window + ExpertStream + PleTable.
// Supports both:
//   1. Prefill (PromptGemm) + decode (Window::run + ExpertStream F_NOCACHE pread + speculative router lookahead)
//   2. Teacher-forced full-vocab (248,320) FP64 KLD, PPL, and top-1 agreement evaluation (--logits-out / --logits-ref)

#include "strata/core/layer.hpp"
#include "strata/core/session.hpp"
#include "strata/kernels/ngram.hpp"
#include "strata/spec/suffix_drafter.hpp"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fcntl.h>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>

namespace strata::metal_engine {
void set_model_shard(const char* shard1);
void* window_create();
void window_destroy(void* w);
bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::SessionState& ss, int max_t, int64_t max_cells,
                 int64_t* n_vocab, std::string& err);
bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0,
                     const std::function<bool(const float*, int64_t, int64_t, std::string&)>& on_chunk, std::string& err);
bool window_run(void* w, int T, const int32_t* tokens, int64_t pos0, bool head, int32_t* out, std::string& err);
bool window_commit(void* w, int n_keep, std::string& err);
bool window_copy_logits(void* w, int t, float* host);
void window_reset_stats(void* w);
}  // namespace strata::metal_engine

static double now_ms() {
    return std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now().time_since_epoch()).count();
}

static double logsumexp_f64(const float* lg, int64_t n, int32_t* argmax_out) {
    float mx = lg[0];
    int32_t am = 0;
    for (int64_t i = 1; i < n; ++i) {
        if (lg[i] > mx) {
            mx = lg[i];
            am = (int32_t) i;
        }
    }
    if (argmax_out) *argmax_out = am;
    double sum = 0.0;
    for (int64_t i = 0; i < n; ++i) {
        sum += std::exp((double) lg[i] - (double) mx);
    }
    return (double) mx + std::log(sum);
}

int main(int argc, char** argv) {
    std::string metal_gguf, ple_gguf, tokens_file, logits_out, logits_ref, logits_ref2;
    int n_expert = 256;
    int64_t max_cells = 4096;
    int n_prompt = 22;
    int n_gen = 64;
    int suffix_draft = 0;

    for (int i = 1; i < argc; ++i) {
        std::string a = argv[i];
        if (a == "--metal-gguf" && i + 1 < argc) metal_gguf = argv[++i];
        else if (a == "--ple-gguf" && i + 1 < argc) ple_gguf = argv[++i];
        else if (a == "--tokens-file" && i + 1 < argc) tokens_file = argv[++i];
        else if (a == "--n-expert" && i + 1 < argc) n_expert = std::atoi(argv[++i]);
        else if (a == "--ctx" && i + 1 < argc) max_cells = std::atoll(argv[++i]);
        else if (a == "--n-prompt" && i + 1 < argc) n_prompt = std::atoi(argv[++i]);
        else if (a == "--n-gen" && i + 1 < argc) n_gen = std::atoi(argv[++i]);
        else if (a == "--suffix-draft" && i + 1 < argc) suffix_draft = std::max(0, std::min(7, std::atoi(argv[++i])));
        else if (a == "--logits-out" && i + 1 < argc) logits_out = argv[++i];
        else if (a == "--logits-ref" && i + 1 < argc) logits_ref = argv[++i];
        else if (a == "--logits-ref2" && i + 1 < argc) logits_ref2 = argv[++i];
    }
    if (metal_gguf.empty()) {
        std::fprintf(stderr, "usage: strata_window_cli --metal-gguf <K256-metal.gguf> [--ple-gguf <shard1.gguf>] "
                             "[--n-expert 256] [--ctx 4096] [--n-prompt 22] [--n-gen 64] [--tokens-file <ids.txt>] "
                             "[--logits-out <out.bin>] [--logits-ref <ref.bin>]\n");
        return 2;
    }

    const bool tf_mode = !logits_out.empty() || !logits_ref.empty();
    std::vector<int64_t> all_ids;
    if (!tokens_file.empty()) {
        std::ifstream f(tokens_file);
        int64_t tid;
        while (f >> tid) all_ids.push_back(tid);
    }
    if (all_ids.empty()) {
        const int64_t def_ids[] = {
            248045, 8678, 198, 2523, 513, 264, 10631, 17313, 13, 248046, 198, 248045, 846, 198, 814, 20139, 1204, 27413, 71284, 3728, 13081, 32907, 12361, 321, 264, 3074, 6700, 11, 321, 948, 264, 28336, 628, 6561, 571, 13, 248046, 198, 248045, 74455, 198, 248068, 198
        };
        const size_t nd = sizeof(def_ids) / sizeof(def_ids[0]);
        for (int i = 0; i < n_prompt + (tf_mode ? n_gen : 0); ++i) all_ids.push_back(def_ids[i % nd]);
    }
    if (tf_mode && (int) all_ids.size() < n_prompt + n_gen) {
        std::fprintf(stderr, "strata_window_cli: teacher-forced mode needs at least n_prompt + n_gen (%d) tokens in --tokens-file, got %zu\n",
                     n_prompt + n_gen, all_ids.size());
        return 1;
    }
    std::vector<int64_t> prompt_ids(all_ids.begin(), all_ids.begin() + std::min<size_t>(all_ids.size(), (size_t) n_prompt));

    const double t_init0 = now_ms();
    std::string err;
    strata::kernels::PleTable ple_table;
    strata::core::SessionState ss{};
    ss.ple_prev[0] = -1;
    ss.ple_prev[1] = -1;
    ss.ple.consts = strata::kernels::ple_artifact_consts();
    if (!ple_gguf.empty()) {
        strata::kernels::PleIoOptions io{};
        io.mode = strata::kernels::PleIo::Direct;
        io.max_inflight = 256;
        io.cache_rows = 1u << 18;
        if (!ple_table.open(ple_gguf, err, io)) {
            std::fprintf(stderr, "strata_window_cli: PleTable::open failed: %s\n", err.c_str());
            return 1;
        }
        ss.ple.table = &ple_table;
    }

    strata::core::ModelGeometry g{};
    g.n_embd = 2560;
    g.hc = 4;
    g.hc_lr = 320;
    g.n_ff = 640;
    g.n_expert = n_expert;
    g.n_layers = 48;

    strata::metal_engine::set_model_shard(metal_gguf.c_str());
    void* w = strata::metal_engine::window_create();
    int64_t n_vocab = 0;
    if (!strata::metal_engine::window_init(w, g, ss, 1 + suffix_draft, max_cells, &n_vocab, err)) {
        std::fprintf(stderr, "strata_window_cli: window_init failed: %s\n", err.c_str());
        return 1;
    }
    const double init_ms = now_ms() - t_init0;

    FILE* fout_lg = nullptr;
    if (!logits_out.empty()) {
        fout_lg = std::fopen(logits_out.c_str(), "wb");
        if (!fout_lg) {
            std::perror("fopen logits_out");
            return 1;
        }
        fcntl(fileno(fout_lg), F_NOCACHE, 1);
    }
    FILE* fref_lg = nullptr;
    if (!logits_ref.empty()) {
        fref_lg = std::fopen(logits_ref.c_str(), "rb");
        if (!fref_lg) {
            std::perror("fopen logits_ref");
            return 1;
        }
        fcntl(fileno(fref_lg), F_NOCACHE, 1);
    }
    FILE* fref2_lg = nullptr;
    if (!logits_ref2.empty()) {
        fref2_lg = std::fopen(logits_ref2.c_str(), "rb");
        if (!fref2_lg) {
            std::perror("fopen logits_ref2");
            return 1;
        }
        fcntl(fileno(fref2_lg), F_NOCACHE, 1);
    }

    strata::spec::SuffixDrafter suffix(3, 32);
    for (int64_t tid : prompt_ids) suffix.append((int32_t) tid);

    // Prefill all but the last prompt token via prompt_run_rows
    const int64_t n_pre = std::max<int64_t>(0, (int64_t) prompt_ids.size() - 1);
    const double t_pre0 = now_ms();
    if (n_pre > 0) {
        if (!strata::metal_engine::prompt_run_rows(prompt_ids.data(), n_pre, 0, nullptr, err)) {
            std::fprintf(stderr, "strata_window_cli: prompt_run_rows failed: %s\n", err.c_str());
            return 1;
        }
    }
    // Run the last prompt token with head=true to get the first output distribution
    int32_t cur_tok = (int32_t) prompt_ids.back();
    int32_t next_tok = 0;
    int64_t pos = n_pre;
    if (!strata::metal_engine::window_run(w, 1, &cur_tok, pos, true, &next_tok, err) ||
        !strata::metal_engine::window_commit(w, 1, err)) {
        std::fprintf(stderr, "strata_window_cli: first window_run failed: %s\n", err.c_str());
        return 1;
    }
    ++pos;
    const double prefill_ms = now_ms() - t_pre0;

    std::vector<float> cur_lg(tf_mode ? (size_t) n_vocab : 0);
    std::vector<float> ref_lg((fref_lg || fref2_lg) ? (size_t) n_vocab : 0);
    double sum_nll_cand = 0.0, sum_nll_ref = 0.0, sum_nll_ref2 = 0.0;
    std::vector<double> klds, klds2;
    int top1_matches = 0, top1_matches2 = 0;

    auto eval_tf_step = [&](int step_idx) -> bool {
        if (!strata::metal_engine::window_copy_logits(w, 0, cur_lg.data())) return false;
        if (fout_lg) {
            if (std::fwrite(cur_lg.data(), sizeof(float), (size_t) n_vocab, fout_lg) != (size_t) n_vocab) return false;
        }
        const int32_t tgt = (int32_t) all_ids[(size_t) n_prompt + (size_t) step_idx];
        int32_t cand_top1 = 0;
        const double lse_cand = logsumexp_f64(cur_lg.data(), n_vocab, &cand_top1);
        sum_nll_cand += -((double) cur_lg[(size_t) tgt] - lse_cand);

        if (fref_lg) {
            if (std::fread(ref_lg.data(), sizeof(float), (size_t) n_vocab, fref_lg) != (size_t) n_vocab) return false;
            int32_t ref_top1 = 0;
            const double lse_ref = logsumexp_f64(ref_lg.data(), n_vocab, &ref_top1);
            sum_nll_ref += -((double) ref_lg[(size_t) tgt] - lse_ref);
            if (cand_top1 == ref_top1) ++top1_matches;
            double kld = 0.0;
            for (int64_t v = 0; v < n_vocab; ++v) {
                const double log_p = (double) ref_lg[(size_t) v] - lse_ref;
                const double log_q = (double) cur_lg[(size_t) v] - lse_cand;
                const double p = std::exp(log_p);
                if (p > 0.0) kld += p * (log_p - log_q);
            }
            klds.push_back(std::max(0.0, kld));
        }
        if (fref2_lg) {
            if (std::fread(ref_lg.data(), sizeof(float), (size_t) n_vocab, fref2_lg) != (size_t) n_vocab) return false;
            int32_t ref2_top1 = 0;
            const double lse_ref2 = logsumexp_f64(ref_lg.data(), n_vocab, &ref2_top1);
            sum_nll_ref2 += -((double) ref_lg[(size_t) tgt] - lse_ref2);
            if (cand_top1 == ref2_top1) ++top1_matches2;
            double kld2 = 0.0;
            for (int64_t v = 0; v < n_vocab; ++v) {
                const double log_p = (double) ref_lg[(size_t) v] - lse_ref2;
                const double log_q = (double) cur_lg[(size_t) v] - lse_cand;
                const double p = std::exp(log_p);
                if (p > 0.0) kld2 += p * (log_p - log_q);
            }
            klds2.push_back(std::max(0.0, kld2));
        }
        return true;
    };

    if (tf_mode && n_gen > 0) {
        if (!eval_tf_step(0)) {
            std::fprintf(stderr, "strata_window_cli: eval_tf_step(0) failed\n");
            return 1;
        }
    }

    std::vector<int32_t> gen_tokens;
    gen_tokens.push_back(next_tok);
    suffix.append(next_tok);
    cur_tok = tf_mode ? (int32_t) all_ids[(size_t) n_prompt] : next_tok;

    strata::metal_engine::window_reset_stats(w);
    const int decode_steps = std::max(0, n_gen - 1);
    const double t_dec0 = now_ms();
    int produced = 0;
    while (produced < decode_steps) {
        int32_t win_in[8] = {cur_tok};
        int32_t win_out[8] = {};
        int nd = 0;
        if (!tf_mode && suffix_draft > 0) {
            const int rem = decode_steps - produced;
            const int max_d = std::min(suffix_draft, std::max(0, rem - 1));
            if (max_d > 0) nd = suffix.propose(max_d, win_in + 1);
        }
        const int T = 1 + nd;
        if (!strata::metal_engine::window_run(w, T, win_in, pos, true, win_out, err)) {
            std::fprintf(stderr, "strata_window_cli: decode step failed: %s\n", err.c_str());
            return 1;
        }
        int n_keep = 1;
        for (int i = 0; i < nd; ++i) {
            if (win_out[i] == win_in[i + 1]) ++n_keep;
            else break;
        }
        if (!strata::metal_engine::window_commit(w, n_keep, err)) {
            std::fprintf(stderr, "strata_window_cli: window_commit failed: %s\n", err.c_str());
            return 1;
        }
        if (tf_mode) {
            if (!eval_tf_step(produced + 1)) {
                std::fprintf(stderr, "strata_window_cli: eval_tf_step(%d) failed\n", produced + 1);
                return 1;
            }
            gen_tokens.push_back(win_out[0]);
            ++produced;
            if (produced < decode_steps) {
                cur_tok = (int32_t) all_ids[(size_t) n_prompt + (size_t) produced];
            }
            pos += 1;
        } else {
            for (int i = 0; i < n_keep && produced < decode_steps; ++i) {
                gen_tokens.push_back(win_out[i]);
                suffix.append(win_out[i]);
                cur_tok = win_out[i];
                ++produced;
            }
            pos += n_keep;
        }
    }
    const double decode_ms = now_ms() - t_dec0;
    const double decode_tok_s = decode_steps > 0 ? (decode_steps * 1000.0 / decode_ms) : 0.0;
    const double decode_ms_tok = decode_steps > 0 ? (decode_ms / decode_steps) : 0.0;
    const double prefill_tok_s = prompt_ids.size() * 1000.0 / prefill_ms;

    if (fout_lg) std::fclose(fout_lg);
    if (fref_lg) std::fclose(fref_lg);
    if (fref2_lg) std::fclose(fref2_lg);

    std::printf("STRATA_RESULT init_ms=%.2f prompt_tokens=%zu prefill_ms=%.2f prefill_tok_s=%.2f "
                "decode_tokens=%d decode_ms=%.2f decode_tok_s=%.2f decode_ms_per_tok=%.2f\n",
                init_ms, prompt_ids.size(), prefill_ms, prefill_tok_s,
                decode_steps, decode_ms, decode_tok_s, decode_ms_tok);
    if (tf_mode && n_gen > 0) {
        const double cand_ppl = std::exp(sum_nll_cand / n_gen);
        if (!klds.empty()) {
            const double ref_ppl = std::exp(sum_nll_ref / n_gen);
            double mean_kld = 0.0;
            for (double x : klds) mean_kld += x;
            mean_kld /= klds.size();
            std::vector<double> sk = klds;
            std::sort(sk.begin(), sk.end());
            const double med_kld = sk[sk.size() / 2];
            const double p90_kld = sk[(size_t) (0.90 * (sk.size() - 1))];
            const double top1_pct = 100.0 * (double) top1_matches / (double) n_gen;
            std::printf("STRATA_KLD eval_tokens=%d ref_ppl=%.4f cand_ppl=%.4f delta_ppl=%.4f mean_kld=%.6f median_kld=%.6f p90_kld=%.6f top1_agree_pct=%.2f\n",
                        n_gen, ref_ppl, cand_ppl, cand_ppl - ref_ppl, mean_kld, med_kld, p90_kld, top1_pct);
            if (!klds2.empty()) {
                const double ref2_ppl = std::exp(sum_nll_ref2 / n_gen);
                double mean_kld2 = 0.0;
                for (double x : klds2) mean_kld2 += x;
                mean_kld2 /= klds2.size();
                std::vector<double> sk2 = klds2;
                std::sort(sk2.begin(), sk2.end());
                const double med_kld2 = sk2[sk2.size() / 2];
                const double p90_kld2 = sk2[(size_t) (0.90 * (sk2.size() - 1))];
                const double top1_pct2 = 100.0 * (double) top1_matches2 / (double) n_gen;
                std::printf("STRATA_KLD_REF2 eval_tokens=%d ref_ppl=%.4f cand_ppl=%.4f delta_ppl=%.4f mean_kld=%.6f median_kld=%.6f p90_kld=%.6f top1_agree_pct=%.2f\n",
                            n_gen, ref2_ppl, cand_ppl, cand_ppl - ref2_ppl, mean_kld2, med_kld2, p90_kld2, top1_pct2);
            }
        } else {
            std::printf("STRATA_PPL eval_tokens=%d cand_ppl=%.4f\n", n_gen, cand_ppl);
        }
    }
    std::printf("GEN_IDS:");
    for (size_t i = 0; i < gen_tokens.size(); ++i) {
        std::printf(" %d", gen_tokens[i]);
    }
    std::printf("\n");
    std::fflush(stdout);

    strata::metal_engine::window_destroy(w);
    return 0;
}
"""
    (ROOT / "metal" / "src" / "strata_window_cli.cpp").write_text(cli_src, encoding="utf-8")

    # Also update convert_k256_q2_0.cpp so SSE search evaluates the fp16-rounded scale d
    conv_path = ROOT / "metal" / "src" / "convert_k256_q2_0.cpp"
    if conv_path.exists():
        conv_txt = conv_path.read_text(encoding="utf-8")
        conv_txt = conv_txt.replace(
            "const float d = amax * f;\n                const float inv = 1.0f / d;",
            "const float d = ref::f16_to_f32(ref::f32_to_f16(amax * f));\n                if (d <= 0.0f) continue;\n                const float inv = 1.0f / d;",
        )
        conv_path.write_text(conv_txt, encoding="utf-8")

    heal_src = Path(__file__).resolve().parent / "heal_k256_q2_0.cpp"
    if heal_src.exists():
        (ROOT / "metal" / "src" / "heal_k256_q2_0.cpp").write_text(heal_src.read_text(encoding="utf-8"), encoding="utf-8")

    cm_path = ROOT / "metal" / "CMakeLists.txt"
    cm_txt = cm_path.read_text(encoding="utf-8")
    if "strata_window_cli" not in cm_txt:
        cm_txt += """
add_executable(strata_window_cli src/strata_window_cli.cpp)
target_link_libraries(strata_window_cli PRIVATE strata_metal_kernels strata_metal_host strata_metal_kernels)
add_executable(convert_k256_q2_0 src/convert_k256_q2_0.cpp)
target_link_libraries(convert_k256_q2_0 PRIVATE mtl_rt)
add_executable(heal_k256_q2_0 src/heal_k256_q2_0.cpp)
target_compile_definitions(heal_k256_q2_0 PRIVATE ACCELERATE_NEW_LAPACK)
target_link_libraries(heal_k256_q2_0 PRIVATE mtl_rt "-framework Accelerate")
"""
        cm_path.write_text(cm_txt, encoding="utf-8")


if __name__ == "__main__":
    reset_git()
    patch_mtl_rt()
    patch_kernels()
    patch_stream()
    patch_window()
    patch_tests_and_route_parity()
    write_strata_window_cli_and_converter()
    print("Patched vendor/strata/metal successfully!")