// Oasis engine: loads an .oasis model and describes RGB images as JSON records, in one C++17 translation
// unit. `python oasis.py build` compiles it into the library that oasis.py loads, as
//   c++ -std=c++17 -O3 -ffast-math -fno-finite-math-only -march=native -fopenmp -fPIC -shared -fvisibility=hidden oasis.cpp -o liboasis.so
//   cl /std:c++17 /O2 /fp:fast /arch:AVX2 /openmp /EHsc /LD oasis.cpp /Fe:oasis.dll
// with -mcpu=native in place of -march=native for Clang on 64-bit Arm. The C API below is all it exports.

#ifndef OASIS_H
#define OASIS_H

#ifdef __cplusplus
extern "C" {
#endif

#if defined(_WIN32)
#define OASIS_API __declspec(dllexport)
#else
#define OASIS_API __attribute__((visibility("default")))
#endif

typedef struct oasis_model oasis_model;

/* Returns NULL on failure; see oasis_last_error(). */
OASIS_API oasis_model* oasis_load(const char* path);
OASIS_API void oasis_free(oasis_model* model);

/* rgb: packed 8-bit RGB, row-major, width * height * 3 bytes.
   options_json: NULL or a JSON object with language (en, de, fr, es, it, nl or pl), keywords (a list
   of words or phrases the description should favour, or one newline-separated string), conf, text_conf,
   max_objects, max_text_lines, max_tokens, masks, mask_format ("largest", "polygons", "merged",
   "bitmap", "none"), mask_threshold, nms, reasoning, sections (any of "scene", "objects", "regions", "text", "relations", "events",
   "reasoning"; text alone skips the decoder), fields (any of "bbox", "score", "mask", "position",
   "description", "count") and form (an ordered object of fields: {"enum": [...]}, {"type": "integer",
   "minimum": a, "maximum": b}, {"type": "string", "maxLength": n} or {"x-source": "bbox", "x-of":
   "<earlier enum field>"}), which returns {"form": ...} instead of a description.
   Writes a NUL-terminated JSON record to out; returns its length, or a negative value on error
   (minus the required size when capacity is too small). */
OASIS_API int oasis_describe(oasis_model* model, const unsigned char* rgb, int width, int height,
                             const char* options_json, char* out, int capacity);

OASIS_API const char* oasis_last_error(void);

#ifdef __cplusplus
}
#endif

#endif

#include <algorithm>
#include <atomic>
#include <cctype>
#include <chrono>
#include <climits>
#include <cmath>
#include <condition_variable>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <memory>
#include <mutex>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>
#include <type_traits>
#include <unordered_map>
#include <utility>
#include <vector>
#ifdef _OPENMP
#include <omp.h>
#endif
#if defined(__F16C__)
#include <immintrin.h>
#endif

// ==== Internals: model container, tensor access, kernel and pipeline declarations

// 64-bit Arm with NEON: the kernels and the image and text stages take their Arm paths.
#if defined(__aarch64__) && defined(__ARM_NEON)
#include <arm_neon.h>
#define OASIS_NEON 1
#endif

namespace oasis {

class SpinPool;

// Tensor types: 0 f32, 1 f16, 2 cq with fp16 group scales, 3 raw bytes, 4 signed integer rows of `bits`
// bits with one f32 scale per row, 5 cq with 8-bit log-coded group scales.
struct Entry {
    int dtype = 0;
    int bits = 0;
    std::vector<int> shape;
    const uint8_t* data = nullptr;
    size_t nbytes = 0;
};

// Quantised matrix [rows = outputs, cols = inputs]: packed indices, then one scale per row and 128-group,
// stored either as fp16 (norms) or as an 8-bit code into a 256-entry table (ncode, nlut).
struct CQ {
    int rows = 0, cols = 0, pc = 0, bits = 0;
    const uint8_t* idx = nullptr;
    const uint16_t* norms = nullptr;
    const uint8_t* ncode = nullptr;
    const float* nlut = nullptr;
    const float* cb = nullptr;
    const uint8_t* zero_rows = nullptr;   // 1 for rows whose every group scale is zero; null when there are none
};

bool inflate_exact(const uint8_t* in, size_t n, uint8_t* out, size_t raw);

struct Model {
    std::vector<uint8_t> file;
    std::vector<std::vector<uint8_t>> inflated;   // tensors stored deflated, expanded at load
    std::unordered_map<std::string, std::string> header;
    std::unordered_map<std::string, Entry> entries;
    std::unordered_map<std::string, std::vector<float>> cache;   // fp16 and integer-row tensors as f32, filled at load
    std::unordered_map<std::string, std::vector<float>> luts;    // decoded 8-bit group-scale codes, filled at load
    std::unordered_map<std::string, std::vector<uint8_t>> zero_rows;   // quantised matrices with all-zero rows
    std::unordered_map<std::string, std::vector<float>> dq;      // dequantised matrices, filled on demand
    std::mutex dq_mutex;
    std::vector<float> codebook;
    int cb_offset[5] = {0, 0, 0, 4, 12};
    int vocab = 0, n_pieces = 0, bos = 2, eos = 1, pad = 0, ref_base = 0;
    std::unordered_map<std::string, int> langs;      // every language tag in the vocabulary
    std::vector<std::string> languages;              // the languages the model describes in
    std::vector<std::string> pieces;
    std::vector<int> piece_types;
    std::unordered_map<std::string, int> piece_id;   // normal pieces, for encoding
    std::vector<std::pair<std::string, int>> markers;   // user-defined pieces, longest first
    int byte_id[256] = {0};                          // byte-fallback pieces
    std::vector<uint32_t> charset;
    std::vector<int> perm1, perm2;
    std::vector<int> seq_order;   // sequency row -> natural row for the 512-point Walsh matrix
    std::shared_ptr<SpinPool> pool;   // decoder workers, created at the first decode
    std::mutex decode_mutex;

    bool load(const std::string& path, std::string& err);
    bool has(const std::string& name) const { return entries.count(name) != 0; }
    const Entry& E(const std::string& name) const;
    const float* F(const std::string& name);      // f32 view, fp16 converted once
    CQ Q(const std::string& name) const;
    const float* Wdq(const std::string& name);    // dequantised [rows, cols], cached
};

// ---- scalar helpers
// The exponent is capped so that no compiler's fast exp can overflow; below -80 the result is ~0 anyway.
inline float sigmoidf(float x) { return 1.0f / (1.0f + std::exp(std::min(-x, 80.0f))); }
inline float siluf(float x) { return x / (1.0f + std::exp(std::min(-x, 80.0f))); }
inline float gelu_tanh(float x) {
    const float k = 0.7978845608028654f;
    return 0.5f * x * (1.0f + std::tanh(k * (x + 0.044715f * x * x * x)));
}
#if defined(OASIS_NEON)
// expf as glibc 2.27 and later computes it on 64-bit Arm, four lanes at a time: x times 32/ln 2 rounded to
// k with ties away from zero, 2^(k/32) from its table, and a cubic in the remainder, in double precision.
// Equal to that library's expf for every float that is not a NaN.
inline float32x2_t exp_lanes(float64x2_t xd) {
    static const uint64_t T[32] = {
        0x3ff0000000000000, 0x3fefd9b0d3158574, 0x3fefb5586cf9890f, 0x3fef9301d0125b51, 0x3fef72b83c7d517b, 0x3fef54873168b9aa,
        0x3fef387a6e756238, 0x3fef1e9df51fdee1, 0x3fef06fe0a31b715, 0x3feef1a7373aa9cb, 0x3feedea64c123422, 0x3feece086061892d,
        0x3feebfdad5362a27, 0x3feeb42b569d4f82, 0x3feeab07dd485429, 0x3feea47eb03a5585, 0x3feea09e667f3bcd, 0x3fee9f75e8ec5f74,
        0x3feea11473eb0187, 0x3feea589994cce13, 0x3feeace5422aa0db, 0x3feeb737b0cdc5e5, 0x3feec49182a3f090, 0x3feed503b23e255d,
        0x3feee89f995ad3ad, 0x3feeff76f2fb5e47, 0x3fef199bdd85529c, 0x3fef3720dcef9069, 0x3fef5818dcfba487, 0x3fef7c97337b9b5f,
        0x3fefa4afa2a490da, 0x3fefd0765b6e4540};
    const float64x2_t z = vmulq_f64(vdupq_n_f64(0x1.71547652b82fep+0 * 32), xd);
    const float64x2_t r = vsubq_f64(z, vrndaq_f64(z));
    const int64x2_t ki = vcvtaq_s64_f64(z);
    const uint64_t k0 = (uint64_t)vgetq_lane_s64(ki, 0), k1 = (uint64_t)vgetq_lane_s64(ki, 1);
    const float64x2_t s = vreinterpretq_f64_u64(vcombine_u64(vcreate_u64(T[k0 % 32] + (k0 << 47)), vcreate_u64(T[k1 % 32] + (k1 << 47))));
    const float64x2_t c = vfmaq_f64(vdupq_n_f64(0x1.ebfce50fac4f3p-3 / 32 / 32), vdupq_n_f64(0x1.c6af84b912394p-5 / 32 / 32 / 32), r);
    const float64x2_t y = vfmaq_f64(vfmaq_f64(vdupq_n_f64(1.0), vdupq_n_f64(0x1.62e42ff0c52d6p-1 / 32), r), c, vmulq_f64(r, r));
    return vcvt_f32_f64(vmulq_f64(y, s));
}
inline float32x4_t expf4(float32x4_t x) {
    const float32x4_t lo = vdupq_n_f32(-0x1.9fe368p6f), hi = vdupq_n_f32(0x1.62e42ep6f);
    const float32x4_t xc = vminq_f32(vmaxq_f32(x, lo), hi);   // NaNs pass through
    const float32x4_t e = vcombine_f32(exp_lanes(vcvt_f64_f32(vget_low_f32(xc))), exp_lanes(vcvt_high_f64_f32(xc)));
    const float32x4_t zero = vbslq_f32(vcltq_f32(x, lo), vdupq_n_f32(0.0f), e);
    return vbslq_f32(vcgtq_f32(x, hi), vdupq_n_f32(INFINITY), zero);
}
// tanhf as glibc 2.41 computes it (CORE-MATH's correctly rounded tanhf), four lanes at a time: a rational
// function in double precision, with its own forms for tiny and large arguments. Equal to the library's
// tanhf for every float.
inline float32x2_t tanh_lanes(float64x2_t z) {
    static const double cn[] = {0x1p+0, 0x1.30877b8b72d33p-3, 0x1.694aa09ae9e5ep-8, 0x1.4101377abb729p-14,
                                0x1.e0392b1db0018p-22, 0x1.2533756e546f7p-30, 0x1.d62e5abe6ae8ap-41, 0x1.b06be534182dep-54};
    static const double cd[] = {0x1p+0, 0x1.ed99131b0ebeap-2, 0x1.0d27ed6c95a69p-5, 0x1.7cbdaca0e9fccp-11,
                                0x1.b4e60b892578ep-18, 0x1.a6f707c5c71abp-26, 0x1.35a8b6e2cd94cp-35, 0x1.ca8230677aa01p-47};
    const float64x2_t z2 = vmulq_f64(z, z), z4 = vmulq_f64(z2, z2), z8 = vmulq_f64(z4, z4);
    auto pair = [&](const double* c, int i) { return vfmaq_f64(vdupq_n_f64(c[i]), z2, vdupq_n_f64(c[i + 1])); };
    float64x2_t n0 = pair(cn, 0), n2 = pair(cn, 2), n4 = pair(cn, 4), n6 = pair(cn, 6);
    n0 = vfmaq_f64(n0, z4, n2);
    n4 = vfmaq_f64(n4, z4, n6);
    n0 = vfmaq_f64(n0, z8, n4);
    float64x2_t d0 = pair(cd, 0), d2 = pair(cd, 2), d4 = pair(cd, 4), d6 = pair(cd, 6);
    d0 = vfmaq_f64(d0, z4, d2);
    d4 = vfmaq_f64(d4, z4, d6);
    d0 = vfmaq_f64(d0, z8, d4);
    return vcvt_f32_f64(vdivq_f64(vmulq_f64(z, n0), d0));
}
inline float32x4_t tanhf4(float32x4_t x) {
    const uint32x4_t ax = vandq_u32(vreinterpretq_u32_f32(x), vdupq_n_u32(0x7fffffff));
    float32x4_t r = vcombine_f32(tanh_lanes(vcvt_f64_f32(vget_low_f32(x))), tanh_lanes(vcvt_high_f64_f32(x)));
    const float32x4_t x2 = vmulq_f32(x, x);
    r = vbslq_f32(vcltq_u32(ax, vdupq_n_u32(115u << 23)), vfmaq_f32(x, x, vmulq_n_f32(x2, -0x1.555556p-2f)), r);
    r = vbslq_f32(vcltq_u32(ax, vdupq_n_u32(102u << 23)), vfmaq_f32(x, vnegq_f32(x), vabsq_f32(x)), r);
    r = vbslq_f32(vceqq_u32(ax, vdupq_n_u32(0)), x, r);
    const uint32x4_t sign = vandq_u32(vreinterpretq_u32_f32(x), vdupq_n_u32(0x80000000));
    const float32x4_t one = vreinterpretq_f32_u32(vorrq_u32(sign, vreinterpretq_u32_f32(vdupq_n_f32(1.0f))));
    const float32x4_t tiny = vreinterpretq_f32_u32(vorrq_u32(sign, vreinterpretq_u32_f32(vdupq_n_f32(0x1p-25f))));
    r = vbslq_f32(vcgtq_u32(ax, vdupq_n_u32(0x41102cb3u)), vsubq_f32(one, tiny), r);
    return vbslq_f32(vcgtq_u32(ax, vdupq_n_u32(0x7f800000u)), vaddq_f32(x, x), r);
}
// o[i] *= sigmoidf(g[i]) and v[i] = siluf(v[i]) for n a multiple of 4, with the same roundings.
inline void mul_sigmoid(float* o, const float* g, int n) {
    const float32x4_t cap = vdupq_n_f32(80.0f), one = vdupq_n_f32(1.0f);
    for (int i = 0; i < n; i += 4) {
        const float32x4_t e = expf4(vminq_f32(vnegq_f32(vld1q_f32(g + i)), cap));
        vst1q_f32(o + i, vmulq_f32(vld1q_f32(o + i), vdivq_f32(one, vaddq_f32(one, e))));
    }
}
inline void silu_inplace(float* v, int n) {
    const float32x4_t cap = vdupq_n_f32(80.0f), one = vdupq_n_f32(1.0f);
    for (int i = 0; i < n; i += 4) {
        const float32x4_t x = vld1q_f32(v + i);
        vst1q_f32(v + i, vdivq_f32(x, vaddq_f32(one, expf4(vminq_f32(vnegq_f32(x), cap)))));
    }
}
#endif
float half_to_float(uint16_t h);
extern bool g_exact_gemv;   // float activations in quantised matrix-vector products (OASIS_EXACT=1)
// Threads for one description: OMP_NUM_THREADS when it is set, otherwise every processor, less one on
// machines with four or fewer, where the rest of the system would otherwise preempt a worker at every step.
int engine_threads();

// Scale of group i (row-major over rows and 128-groups) of a quantised matrix.
inline float cq_scale(const CQ& q, size_t i) { return q.ncode ? q.nlut[q.ncode[i]] : half_to_float(q.norms[i]); }

// Activations prepared once for several quantised matrix-vector products with the same input:
// Walsh-rotated per 128-group, int8 per group, reordered to the packed-index layout.
struct QX {
    std::vector<int8_t> xs, xa;
    std::vector<float> sx;
    int pc = 0, bits = 0;
};

// Codebook entries as int8 with one scale, shared by every row of a matrix of the same width.
struct CbTable {
    alignas(16) int8_t tbl[16];
    float cs;
};

// ---- kernels
void prep_x(const float* x, int cols, int bits, QX& out);
CbTable cb_table(const CQ& q);
float cq_row(const CQ& q, const QX& qx, int r, const CbTable& t);   // one output of W x
void cq_gemv_q(const CQ& q, const QX& qx, float* y);
void cq_gemm_q(const CQ& q, const float* A, int M, float* C);   // int8-activation matrix product
struct Model;
void qmm(Model& m, const std::string& name, const float* A, int M, float* C);   // C = A W^T for a quantised W
float dotf(const float* a, const float* b, int n);
void wht_inplace(float* x, int n);                         // natural order, normalised
void wht_seq512(const Model& m, const float* x, float* out);   // sequency-ordered 512-point transform
void cq_dequant_row(const CQ& q, int row, float* out);
void cq_gemv(const CQ& q, const float* x, float* y);        // y[rows] = W x
void gemm_nt(const float* A, int M, int K, const float* W, int N, float* C);   // C = A W^T, W [N, K]
void gemm_nn(const float* A, int M, int K, const float* B, int N, float* C);   // C = A B,   B [K, N]
void zc_norm(const float* x, int n, const float* scale, float* out);
void rms_unit(const float* x, int n, float* out);
void band65(const float* u, const float* w, int n, float* out);   // per-output 65-tap band, zero edges

// ---- pipeline
struct Box {
    float x0, y0, x1, y1;  // canvas-normalised [0, 1]
    float score;
    int ci, cj;            // centre cell on the 64 x 64 grid
};

struct Meta {
    int W = 0, H = 0, w = 0, h = 0;
    float r = 1.0f;
};

struct Encoded {
    Meta meta;
    std::vector<float> E;    // [256, 512]
    std::vector<float> D;    // [64, 64, 128]
    std::vector<float> S4;   // [128, 128, 64]
    std::vector<uint8_t> valid32, valid16, valid8, valid4;
    std::vector<float> det;  // [64, 64, 10]
};

// One field of a form request, in request order.
struct FormField {
    enum Kind { CHOICE, COUNT, TEXT, BOX } kind = CHOICE;
    std::string name;
    std::vector<std::string> choices;   // CHOICE: candidate texts
    std::vector<std::string> values;    // CHOICE: each choice as JSON (string or number)
    long lo = 0, hi = 0;                // COUNT: inclusive range
    int max_length = 60;                // TEXT
    std::string of;                     // BOX: the earlier choice field it locates
};

struct Options {
    std::string language = "en";
    float conf = 0.3f, text_conf = 0.35f, mask_threshold = 0.5f;
    int max_objects = 20, max_text_lines = 20, max_tokens = 512;
    std::string mask_format = "largest";
    bool masks = true, nms = true, reasoning = true;
    std::vector<std::string> sections, fields, keywords;
    std::vector<FormField> form;
};

Encoded encode(Model& m, const uint8_t* rgb, int W, int H);
void detect(const Encoded& enc, float conf, float text_conf, bool nms, int max_obj, int max_txt,
            std::vector<Box>& objects, std::vector<Box>& texts);
std::vector<float> mask_prob(Model& m, const Encoded& enc, const Box& b, std::vector<float>* protos);
std::string read_text(Model& m, const uint8_t* rgb, int W, int H, float x0, float y0, float x1, float y1);
void memory(Model& m, const Encoded& enc, const std::vector<Box>& objects, const std::vector<Box>& texts,
            std::vector<float>& M, std::vector<uint8_t>& mvalid, std::vector<int>& refs);
// keywords: words and phrases whose tokens are favoured while the description is decoded.
std::vector<int> generate(Model& m, const std::vector<float>& M, const std::vector<uint8_t>& mvalid,
                          const std::vector<int>& refs, const std::string& language, int max_tokens,
                          const std::vector<std::string>& keywords = {});
std::string detokenize(const Model& m, const std::vector<int>& ids);
std::vector<int> encode_text(const Model& m, const std::string& text);   // SentencePiece BPE, no dummy prefix
// Fills a form against the decoder memory; ref_boxes are the 0-1000 boxes of refs. Returns the "form" object.
std::string fill_form(Model& m, const std::vector<float>& M, const std::vector<uint8_t>& mvalid, const std::vector<int>& refs,
                      const std::vector<std::vector<int>>& ref_boxes, const std::string& language, const std::vector<FormField>& form);
std::string describe_json(Model& m, const uint8_t* rgb, int W, int H, const Options& o, double* timings);

}  // namespace oasis

// ==== Worker threads that stay awake for the length of one decode

//
// Each decoded token runs dozens of parallel steps that last a few microseconds. Workers that an
// OpenMP runtime parks between such steps take longer to wake than the steps take to run, so the
// decoder shares its work through this pool instead: while a decode is in progress the workers spin
// on a generation counter, and outside it they sleep on a condition variable.

#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || defined(_M_IX86)
#include <immintrin.h>
#define OASIS_PAUSE() _mm_pause()
#elif defined(__aarch64__) && (defined(__GNUC__) || defined(__clang__))
#define OASIS_PAUSE() __asm__ __volatile__("yield" ::: "memory")
#else
#define OASIS_PAUSE() std::this_thread::yield()
#endif

namespace oasis {

class SpinPool {
public:
    explicit SpinPool(int threads) : n_(std::max(1, threads)) {
        for (int i = 1; i < n_; i++) workers_.emplace_back([this] { loop(); });
    }

    ~SpinPool() {
        {
            std::lock_guard<std::mutex> lock(mutex_);
            stop_.store(true);
            awake_.store(true);
        }
        cv_.notify_all();
        for (auto& t : workers_) t.join();
    }

    SpinPool(const SpinPool&) = delete;
    SpinPool& operator=(const SpinPool&) = delete;

    int size() const { return n_; }

    void wake() {
        {
            std::lock_guard<std::mutex> lock(mutex_);
            awake_.store(true);
        }
        cv_.notify_all();
    }

    // Workers notice within a few hundred spins and go back to sleep.
    void sleep() {
        std::lock_guard<std::mutex> lock(mutex_);
        awake_.store(false);
    }

    // Calls f(i) once for every i in [0, n). Indices are claimed in chunks by whichever thread is free,
    // and each is computed entirely by one thread, so results do not depend on the thread count. The call
    // returns once every index is done: a worker the system has not scheduled meanwhile takes no part in
    // the step instead of holding the others up.
    template <class F>
    void run(int n, int chunk, F&& f) {
        if (n <= 0) return;
        if (n_ == 1 || n <= chunk) {
            for (int i = 0; i < n; i++) f(i);
            return;
        }
        if (!awake_.load(std::memory_order_relaxed)) wake();
        using G = typename std::remove_reference<F>::type;
        const uint64_t gen = (ticket_.load(std::memory_order_relaxed) >> 32) + 1;
        ticket_.store(gen << 32 | CLOSED, std::memory_order_relaxed);
        fn_.store([](void* ctx, int b, int e) {
            G& g = *static_cast<G*>(ctx);
            for (int i = b; i < e; i++) g(i);
        }, std::memory_order_relaxed);
        ctx_.store((void*)&f, std::memory_order_relaxed);
        chunk_.store(std::max(1, chunk), std::memory_order_relaxed);
        done_.store(0, std::memory_order_relaxed);
        items_.store(n, std::memory_order_release);
        ticket_.store(gen << 32, std::memory_order_release);
        work(gen);
        while (done_.load(std::memory_order_acquire) < n) OASIS_PAUSE();
    }

private:
    // The ticket holds the step's generation in its high half and the next unclaimed index in its low half.
    // run() closes the new generation before it writes the step and publishes the step's size after that, so
    // a worker holding a ticket of an earlier step fails its claim: either it reads that step's size, which
    // its index has passed, or it reads the new size, and its claim then meets the closed ticket. A step
    // cannot finish while a claimed chunk is outstanding, so the step's function and context stay valid for
    // every successful claim.
    static constexpr uint32_t CLOSED = 0xffffffffu;

    void work(uint64_t gen) {
        for (;;) {
            uint64_t t = ticket_.load(std::memory_order_acquire);
            if ((t >> 32) != gen) return;
            if ((uint32_t)t == CLOSED) {
                OASIS_PAUSE();
                continue;
            }
            const int items = items_.load(std::memory_order_acquire), chunk = chunk_.load(std::memory_order_relaxed);
            const int b = (int)(uint32_t)t;
            if (b >= items) return;
            if (!ticket_.compare_exchange_weak(t, t + (uint64_t)chunk, std::memory_order_acq_rel, std::memory_order_relaxed)) continue;
            const int e = std::min(items, b + chunk);
            fn_.load(std::memory_order_relaxed)(ctx_.load(std::memory_order_relaxed), b, e);
            done_.fetch_add(e - b, std::memory_order_release);
        }
    }

    void loop() {
        uint64_t seen = ticket_.load(std::memory_order_acquire) >> 32;
        for (;;) {
            uint64_t g;
            int spins = 0;
            while ((g = ticket_.load(std::memory_order_acquire) >> 32) == seen) {
                if (stop_.load(std::memory_order_relaxed)) return;
                if (++spins >= 256) {
                    spins = 0;
                    if (!awake_.load(std::memory_order_relaxed)) {
                        std::unique_lock<std::mutex> lock(mutex_);
                        cv_.wait(lock, [&] {
                            return awake_.load() || stop_.load() || (ticket_.load(std::memory_order_acquire) >> 32) != seen;
                        });
                    }
                }
                OASIS_PAUSE();
            }
            seen = g;
            work(g);
        }
    }

    int n_;
    std::vector<std::thread> workers_;
    std::mutex mutex_;
    std::condition_variable cv_;
    std::atomic<bool> awake_{false}, stop_{false};
    std::atomic<uint64_t> ticket_{0};
    std::atomic<int> items_{0}, chunk_{1}, done_{0};
    std::atomic<void (*)(void*, int, int)> fn_{nullptr};
    std::atomic<void*> ctx_{nullptr};
};

}  // namespace oasis

// ==== zlib inflate (RFC 1950 and 1951), for the tensors stored deflated

namespace oasis {
namespace {

// Canonical Huffman code. Codes of up to 9 bits resolve in one lookup of the next 9 input bits, longer
// ones by walking the code lengths.
struct Huffman {
    uint16_t fast[512];   // length << 9 | symbol, or 0 when the next 9 bits start no code of up to 9 bits
    uint16_t count[16];   // codes of each length
    uint16_t symbol[288]; // symbols in code order

    bool build(const uint8_t* lengths, int n) {
        std::memset(count, 0, sizeof count);
        for (int i = 0; i < n; i++) count[lengths[i]]++;
        count[0] = 0;
        int left = 1;
        for (int l = 1; l < 16; l++) {
            left = 2 * left - count[l];
            if (left < 0) return false;   // over-subscribed
        }
        uint16_t next[16] = {0};
        for (int l = 1; l < 15; l++) next[l + 1] = (uint16_t)(next[l] + count[l]);
        for (int i = 0; i < n; i++)
            if (lengths[i]) symbol[next[lengths[i]]++] = (uint16_t)i;
        std::memset(fast, 0, sizeof fast);
        for (int l = 1, code = 0, k = 0; l <= 9; l++, code <<= 1)
            for (int j = 0; j < count[l]; j++, code++, k++) {
                int rev = 0;   // deflate sends a code's bits from its most significant one
                for (int b = 0; b < l; b++) rev |= ((code >> b) & 1) << (l - 1 - b);
                for (int f = rev; f < 512; f += 1 << l) fast[f] = (uint16_t)(l << 9 | symbol[k]);
            }
        return true;
    }
};

// Bits in stream order, least significant first. Past the end it supplies zero bytes and counts them, so
// that a stream which ends early is caught by overrun().
struct BitReader {
    const uint8_t* p;
    const uint8_t* end;
    uint64_t buf = 0;
    int cnt = 0;
    size_t past = 0;

    void fill() {
        while (cnt <= 56) {
            uint64_t byte = 0;
            if (p < end) byte = *p++;
            else past++;
            buf |= byte << cnt;
            cnt += 8;
        }
    }
    uint32_t bits(int n) {
        if (cnt < n) fill();
        uint32_t v = (uint32_t)(buf & ((1ull << n) - 1));
        buf >>= n;
        cnt -= n;
        return v;
    }
    bool overrun() const { return past * 8 > (size_t)cnt; }
};

int decode(BitReader& br, const Huffman& h) {
    if (br.cnt < 16) br.fill();
    int e = h.fast[br.buf & 511];
    if (e) {
        br.buf >>= e >> 9;
        br.cnt -= e >> 9;
        return e & 511;
    }
    int code = 0, first = 0, index = 0;
    for (int l = 1; l < 16; l++) {
        code |= (int)(br.buf & 1);
        br.buf >>= 1;
        br.cnt--;
        int c = h.count[l];
        if (code - first < c) return h.symbol[index + code - first];
        index += c;
        first = (first + c) << 1;
        code <<= 1;
    }
    return -1;
}

// The literal/length and distance codes of a block with dynamic Huffman codes.
bool read_tables(BitReader& br, Huffman& lit, Huffman& dist) {
    static const uint8_t ORDER[19] = {16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15};
    int nlit = (int)br.bits(5) + 257, ndist = (int)br.bits(5) + 1, ncode = (int)br.bits(4) + 4;
    if (nlit > 286 || ndist > 30) return false;
    uint8_t cl[19] = {0};
    for (int i = 0; i < ncode; i++) cl[ORDER[i]] = (uint8_t)br.bits(3);
    Huffman code;
    if (!code.build(cl, 19)) return false;
    uint8_t lengths[286 + 30];
    for (int k = 0; k < nlit + ndist;) {
        int s = decode(br, code);
        if (s < 0) return false;
        if (s < 16) {
            lengths[k++] = (uint8_t)s;
            continue;
        }
        int rep, value = 0;
        if (s == 16) {
            if (k == 0) return false;
            value = lengths[k - 1];
            rep = 3 + (int)br.bits(2);
        } else if (s == 17) {
            rep = 3 + (int)br.bits(3);
        } else {
            rep = 11 + (int)br.bits(7);
        }
        if (k + rep > nlit + ndist) return false;
        while (rep--) lengths[k++] = (uint8_t)value;
    }
    return lengths[256] != 0 && lit.build(lengths, nlit) && dist.build(lengths + nlit, ndist);
}

const uint16_t LEN_BASE[29] = {3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115,
                               131, 163, 195, 227, 258};
const uint8_t LEN_EXTRA[29] = {0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0};
const uint16_t DIST_BASE[30] = {1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537,
                                2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577};
const uint8_t DIST_EXTRA[30] = {0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13};

uint32_t adler32(const uint8_t* p, size_t n) {
    uint32_t a = 1, b = 0;
    while (n) {
        size_t k = std::min<size_t>(n, 5552);   // the most bytes before b can overflow
        n -= k;
        while (k--) {
            a += *p++;
            b += a;
        }
        a %= 65521;
        b %= 65521;
    }
    return b << 16 | a;
}

}  // namespace

// Inflates a zlib stream into exactly raw bytes; false when the stream is damaged, fails its checksum or
// holds another length.
bool inflate_exact(const uint8_t* in, size_t n, uint8_t* out, size_t raw) {
    if (n < 6 || (in[0] & 15) != 8 || (in[0] >> 4) > 7 || ((in[0] << 8) | in[1]) % 31 || (in[1] & 32)) return false;
    static const std::pair<Huffman, Huffman> fixed = [] {
        uint8_t l[288 + 30];
        for (int i = 0; i < 288; i++) l[i] = i < 144 ? 8 : i < 256 ? 9 : i < 280 ? 7 : 8;
        for (int i = 0; i < 30; i++) l[288 + i] = 5;
        std::pair<Huffman, Huffman> t;
        t.first.build(l, 288);
        t.second.build(l + 288, 30);
        return t;
    }();
    BitReader br{in + 2, in + n - 4};
    Huffman lit, dist;
    size_t o = 0;
    for (bool last = false; !last;) {
        last = br.bits(1) != 0;
        int type = (int)br.bits(2);
        if (type == 0) {
            int drop = br.cnt & 7;   // stored blocks start on a byte boundary
            br.buf >>= drop;
            br.cnt -= drop;
            uint32_t len = br.bits(16), nlen = br.bits(16);
            if ((len ^ 0xffff) != nlen || len > raw - o) return false;
            for (uint32_t i = 0; i < len; i++) out[o++] = (uint8_t)br.bits(8);
        } else if (type == 1 || type == 2) {
            const Huffman* L = &fixed.first;
            const Huffman* D = &fixed.second;
            if (type == 2) {
                if (!read_tables(br, lit, dist)) return false;
                L = &lit;
                D = &dist;
            }
            for (;;) {
                int s = decode(br, *L);
                if (s < 256) {
                    if (s < 0 || o == raw) return false;
                    out[o++] = (uint8_t)s;
                } else if (s == 256) {
                    break;
                } else {
                    s -= 257;
                    if (s >= 29) return false;
                    size_t len = LEN_BASE[s] + br.bits(LEN_EXTRA[s]);
                    int d = decode(br, *D);
                    if (d < 0 || d >= 30) return false;
                    size_t back = DIST_BASE[d] + br.bits(DIST_EXTRA[d]);
                    if (back > o || len > raw - o) return false;
                    for (size_t i = 0; i < len; i++, o++) out[o] = out[o - back];
                }
            }
        } else {
            return false;
        }
        if (br.overrun()) return false;
    }
    const uint8_t* t = in + n - 4;
    uint32_t want = (uint32_t)t[0] << 24 | (uint32_t)t[1] << 16 | (uint32_t)t[2] << 8 | t[3];
    return o == raw && adler32(out, raw) == want;
}

}  // namespace oasis

// ==== .oasis container loading and tensor access

namespace oasis {

float half_to_float(uint16_t h) {
#if defined(__F16C__)
    return _cvtsh_ss(h);
#elif defined(__aarch64__) && (defined(__GNUC__) || defined(__clang__))
    __fp16 v;
    std::memcpy(&v, &h, 2);
    return (float)v;
#else
    uint32_t sign = (uint32_t)(h & 0x8000) << 16;
    uint32_t exp = (h >> 10) & 0x1f;
    uint32_t man = h & 0x3ff;
    uint32_t bits;
    if (exp == 0) {
        if (man == 0) {
            bits = sign;
        } else {
            exp = 127 - 15 + 1;
            while (!(man & 0x400)) {
                man <<= 1;
                exp--;
            }
            man &= 0x3ff;
            bits = sign | (exp << 23) | (man << 13);
        }
    } else if (exp == 31) {
        bits = sign | 0x7f800000 | (man << 13);
    } else {
        bits = sign | ((exp + 127 - 15) << 23) | (man << 13);
    }
    float f;
    std::memcpy(&f, &bits, 4);
    return f;
#endif
}

template <typename T>
static T rd(const uint8_t*& p) {
    T v;
    std::memcpy(&v, p, sizeof(T));
    p += sizeof(T);
    return v;
}

bool Model::load(const std::string& path, std::string& err) {
    std::ifstream f(path, std::ios::binary);
    if (!f) {
        err = "cannot open " + path;
        return false;
    }
    file.assign(std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>());
    if (file.size() < 16 || std::memcmp(file.data(), "OASIS\0v1", 8) != 0) {
        err = "not an .oasis v1 file";
        return false;
    }
    const uint8_t* p = file.data() + 8;
    uint32_t hl = rd<uint32_t>(p);
    std::string head((const char*)p, hl);
    p += hl;
    std::istringstream hs(head);
    std::string line;
    while (std::getline(hs, line)) {
        auto sp = line.find(' ');
        if (sp != std::string::npos) header[line.substr(0, sp)] = line.substr(sp + 1);
    }
    uint32_t n = rd<uint32_t>(p);
    struct Raw { std::string name; Entry e; uint64_t off; int flags; };
    std::vector<Raw> raws;
    for (uint32_t i = 0; i < n; i++) {
        uint16_t nl = rd<uint16_t>(p);
        Raw r;
        r.name.assign((const char*)p, nl);
        p += nl;
        r.e.dtype = rd<uint8_t>(p);
        r.e.bits = rd<uint8_t>(p);
        int nd = rd<uint8_t>(p);
        r.flags = rd<uint8_t>(p);
        for (int d = 0; d < nd; d++) r.e.shape.push_back((int)rd<uint32_t>(p));
        r.off = rd<uint64_t>(p);
        r.e.nbytes = (size_t)rd<uint64_t>(p);
        raws.push_back(std::move(r));
    }
    size_t data_start = ((size_t)(p - file.data()) + 63) / 64 * 64;
    for (auto& r : raws) {
        if (data_start + r.off + r.e.nbytes > file.size()) {
            err = "truncated tensor " + r.name;
            return false;
        }
        r.e.data = file.data() + data_start + r.off;
        if (r.flags & 1) {
            // Deflated tensor: u64 size of the stored bytes, then a zlib stream of them.
            if (r.e.nbytes < 8) { err = "bad compressed tensor " + r.name; return false; }
            uint64_t raw;
            std::memcpy(&raw, r.e.data, 8);
            inflated.emplace_back(raw);
            if (!inflate_exact(r.e.data + 8, r.e.nbytes - 8, inflated.back().data(), raw)) {
                err = "cannot inflate tensor " + r.name;
                return false;
            }
            r.e.data = inflated.back().data();
            r.e.nbytes = raw;
        }
        entries[r.name] = r.e;
    }
    auto num = [&](const char* k) { return std::stoi(header.at(k)); };
    vocab = num("vocab");
    n_pieces = num("n_pieces");
    bos = num("bos");
    eos = num("eos");
    pad = num("pad");
    ref_base = num("ref_base");
    std::istringstream ls(header.at("langs"));
    std::string tokn;
    std::vector<std::string> all;
    while (ls >> tokn) {
        auto c = tokn.find(':');
        langs[tokn.substr(0, c)] = std::stoi(tokn.substr(c + 1));
        all.push_back(tokn.substr(0, c));
    }
    auto lang_line = header.find("languages");
    if (lang_line != header.end()) {
        std::istringstream as(lang_line->second);
        while (as >> tokn)
            if (langs.count(tokn)) languages.push_back(tokn);
    }
    if (languages.empty()) languages = all;
    const float* cbp = (const float*)E("quant/codebook").data;
    codebook.assign(cbp, cbp + E("quant/codebook").shape[0]);
    const Entry& pe = E("tokenizer/pieces");
    const uint8_t* q = pe.data;
    uint32_t np = rd<uint32_t>(q);
    for (uint32_t i = 0; i < np; i++) {
        int t = rd<uint8_t>(q);
        uint16_t l = rd<uint16_t>(q);
        pieces.emplace_back((const char*)q, l);
        piece_types.push_back(t);
        q += l;
        const std::string& p = pieces.back();
        if (t == 0) piece_id.emplace(p, (int)i);
        else if (t == 3 && !p.empty()) markers.push_back({p, (int)i});
        else if (t == 4 && p.size() == 6 && p.compare(0, 3, "<0x") == 0) byte_id[std::stoi(p.substr(3, 2), nullptr, 16) & 255] = (int)i;
    }
    std::stable_sort(markers.begin(), markers.end(), [](const auto& a, const auto& b) { return a.first.size() > b.first.size(); });
    const Entry& ce = E("reader/charset");
    for (int i = 0; i < ce.shape[0]; i++) charset.push_back((uint32_t)((const float*)ce.data)[i]);
    for (int i = 0; i < 512; i++) {
        perm1.push_back((int)((const float*)E("hmlp/perm1").data)[i]);
        perm2.push_back((int)((const float*)E("hmlp/perm2").data)[i]);
    }
    seq_order.resize(512);
    for (int j = 0; j < 512; j++) {
        int g = j ^ (j >> 1), s = 0;
        for (int b = 0; b < 9; b++)
            if (g & (1 << b)) s |= 1 << (8 - b);
        seq_order[j] = s;
    }
    // Expand every fp16 and integer-row tensor now, and decode every 8-bit scale table, so that lookups
    // during inference never write to the model.
    for (auto& kv : entries) {
        const Entry& e = kv.second;
        if (e.dtype == 1) {
            size_t n = e.nbytes / 2;
            std::vector<float> v(n);
            const uint16_t* h = (const uint16_t*)e.data;
            for (size_t i = 0; i < n; i++) v[i] = half_to_float(h[i]);
            cache[kv.first] = std::move(v);
        } else if (e.dtype == 4) {
            // Signed integers of e.bits bits in one LSB-first bit stream, then one f32 scale per row.
            int rows = e.shape[0], b = e.bits;
            size_t cols = 1;
            for (size_t d = 1; d < e.shape.size(); d++) cols *= (size_t)e.shape[d];
            size_t n = (size_t)rows * cols, packed = (n * b + 7) / 8;
            if (b < 2 || b > 16 || e.nbytes != packed + (size_t)rows * 4) { err = "bad integer tensor " + kv.first; return false; }
            std::vector<float> v(n);
            for (int r = 0; r < rows; r++) {
                float s;
                std::memcpy(&s, e.data + packed + (size_t)r * 4, 4);
                for (size_t c = 0; c < cols; c++) {
                    size_t bit = ((size_t)r * cols + c) * b;
                    uint32_t w = 0;
                    for (int k = 0; k < b; k++, bit++) w |= (uint32_t)(e.data[bit >> 3] >> (bit & 7) & 1) << k;
                    int32_t x = (int32_t)(w << (32 - b)) >> (32 - b);
                    v[(size_t)r * cols + c] = s * (float)x;
                }
            }
            cache[kv.first] = std::move(v);
        } else if (e.dtype == 5) {
            // Code 0 is a zero scale; codes 1..255 step evenly in log scale from exp(lo) to exp(hi).
            float lo, hi;
            std::memcpy(&lo, e.data + e.nbytes - 8, 4);
            std::memcpy(&hi, e.data + e.nbytes - 4, 4);
            std::vector<float> lut(256, 0.0f);
            for (int c = 1; c < 256; c++) lut[c] = std::exp(lo + (hi - lo) * (float)(c - 1) / 254.0f);
            luts[kv.first] = std::move(lut);
        }
        if (e.dtype == 2 || e.dtype == 5) {
            // A row whose every group scale is zero multiplies to zero; the kernels skip it.
            int rows = e.shape[0];
            for (size_t d = 1; d + 1 < e.shape.size(); d++) rows *= e.shape[d];
            const int groups = (e.shape.back() + 127) / 128;
            const uint8_t* sc = e.data + (size_t)rows * groups * 128 * e.bits / 8;
            std::vector<uint8_t> z(rows, 1);
            bool any = false;
            for (int r = 0; r < rows; r++) {
                for (int g = 0; g < groups && z[r]; g++) {
                    size_t i = (size_t)r * groups + g;
                    z[r] = e.dtype == 5 ? sc[i] == 0 : (((const uint16_t*)sc)[i] & 0x7fff) == 0;
                }
                any = any || z[r];
            }
            if (any) zero_rows[kv.first] = std::move(z);
        }
    }
#if defined(OASIS_NEON)
    // On Arm the HConvGLU bands are kept tap-major ([65, rows]), the order the band kernel reads them in.
    for (auto& kv : cache) {
        const std::string& k = kv.first;
        auto ends = [&](const char* s) { size_t n = std::strlen(s); return k.size() >= n && k.compare(k.size() - n, n, s) == 0; };
        if (!ends("hconv/gate_conv") && !ends("hconv/value_conv") && !ends("hconv/out_conv")) continue;
        const Entry& e = entries.at(k);
        if (e.shape.size() != 2 || e.shape[0] != 512 || e.shape[1] != 65) continue;
        std::vector<float> t((size_t)65 * 512);
        for (int r = 0; r < 512; r++)
            for (int j = 0; j < 65; j++) t[(size_t)j * 512 + r] = kv.second[(size_t)r * 65 + j];
        kv.second.swap(t);
    }
#endif
    return true;
}

const Entry& Model::E(const std::string& name) const {
    auto it = entries.find(name);
    if (it == entries.end()) throw std::runtime_error("missing tensor " + name);
    return it->second;
}

const float* Model::F(const std::string& name) {
    const Entry& e = E(name);
    if (e.dtype == 0) return (const float*)e.data;
    if (e.dtype != 1 && e.dtype != 4) throw std::runtime_error("tensor is not a float array: " + name);
    return cache.at(name).data();
}

CQ Model::Q(const std::string& name) const {
    const Entry& e = E(name);
    if (e.dtype != 2 && e.dtype != 5) throw std::runtime_error("tensor is not quantised: " + name);
    CQ q;
    q.rows = e.shape[0];
    for (size_t d = 1; d + 1 < e.shape.size(); d++) q.rows *= e.shape[d];
    q.cols = e.shape.back();
    q.pc = (q.cols + 127) / 128 * 128;
    q.bits = e.bits;
    q.idx = e.data;
    const uint8_t* scales = e.data + (size_t)q.rows * q.pc * q.bits / 8;
    if (e.dtype == 5) {
        q.ncode = scales;
        q.nlut = luts.at(name).data();
    } else {
        q.norms = (const uint16_t*)scales;
    }
    q.cb = codebook.data() + cb_offset[q.bits];
    auto zr = zero_rows.find(name);
    if (zr != zero_rows.end()) q.zero_rows = zr->second.data();
    return q;
}

const float* Model::Wdq(const std::string& name) {
    std::lock_guard<std::mutex> lock(dq_mutex);
    std::string key = "#dq:" + name;
    auto c = dq.find(key);
    if (c != dq.end()) return c->second.data();
    CQ q = Q(name);
    std::vector<float> v((size_t)q.rows * q.cols);
    for (int r = 0; r < q.rows; r++) cq_dequant_row(q, r, v.data() + (size_t)r * q.cols);
    return (dq[key] = std::move(v)).data();
}

}  // namespace oasis

// ==== Numeric kernels: Walsh-Hadamard transforms, codebook matrices, GEMM, norms and bands

namespace oasis {

int engine_threads() {
    static const int n = [] {
#ifdef _OPENMP
        const char* e = std::getenv("OMP_NUM_THREADS");
        if (e && *e) return std::max(1, omp_get_max_threads());
        const int p = omp_get_num_procs();
#else
        const int p = (int)std::thread::hardware_concurrency();
#endif
        return p <= 4 ? std::max(1, p - 1) : p;
    }();
    return n;
}

void wht_inplace(float* x, int n) {
#if defined(OASIS_NEON)
    // The same sums and differences on every element as the loops below, four elements at a time.
    if (n >= 8) {
        for (int i = 0; i < n; i += 8) {   // pairs one apart
            const float32x4_t a = vld1q_f32(x + i), b = vld1q_f32(x + i + 4);
            const float32x4_t e = vuzp1q_f32(a, b), o = vuzp2q_f32(a, b);
            const float32x4_t s = vaddq_f32(e, o), d = vsubq_f32(e, o);
            vst1q_f32(x + i, vzip1q_f32(s, d));
            vst1q_f32(x + i + 4, vzip2q_f32(s, d));
        }
        for (int i = 0; i < n; i += 8) {   // pairs two apart
            const float64x2_t a = vreinterpretq_f64_f32(vld1q_f32(x + i)), b = vreinterpretq_f64_f32(vld1q_f32(x + i + 4));
            const float32x4_t lo = vreinterpretq_f32_f64(vzip1q_f64(a, b)), hi = vreinterpretq_f32_f64(vzip2q_f64(a, b));
            const float64x2_t s = vreinterpretq_f64_f32(vaddq_f32(lo, hi)), d = vreinterpretq_f64_f32(vsubq_f32(lo, hi));
            vst1q_f32(x + i, vreinterpretq_f32_f64(vzip1q_f64(s, d)));
            vst1q_f32(x + i + 4, vreinterpretq_f32_f64(vzip2q_f64(s, d)));
        }
        for (int len = 4; len < n; len <<= 1)
            for (int i = 0; i < n; i += len << 1)
                for (int j = i; j < i + len; j += 4) {
                    const float32x4_t a = vld1q_f32(x + j), b = vld1q_f32(x + j + len);
                    vst1q_f32(x + j, vaddq_f32(a, b));
                    vst1q_f32(x + j + len, vsubq_f32(a, b));
                }
        const float s = 1.0f / std::sqrt((float)n);
        for (int i = 0; i < n; i += 4) vst1q_f32(x + i, vmulq_n_f32(vld1q_f32(x + i), s));
        return;
    }
#endif
    for (int len = 1; len < n; len <<= 1)
        for (int i = 0; i < n; i += len << 1)
            for (int j = i; j < i + len; j++) {
                float a = x[j], b = x[j + len];
                x[j] = a + b;
                x[j + len] = a - b;
            }
    float s = 1.0f / std::sqrt((float)n);
    for (int i = 0; i < n; i++) x[i] *= s;
}

void wht_seq512(const Model& m, const float* x, float* out) {
    float t[512];
    std::memcpy(t, x, sizeof t);
    wht_inplace(t, 512);
    for (int j = 0; j < 512; j++) out[j] = t[m.seq_order[j]];
}

static inline int cq_index(const uint8_t* base, size_t pos, int bits) {
    size_t bit = pos * bits;
    size_t byte = bit >> 3;
    int shift = (int)(bit & 7);
    unsigned v = base[byte] >> shift;
    if (shift + bits > 8) v |= (unsigned)base[byte + 1] << (8 - shift);
    return (int)(v & ((1u << bits) - 1));
}

void cq_dequant_row(const CQ& q, int row, float* out) {
    int groups = q.pc / 128;
    float g[128];
    for (int gi = 0; gi < groups; gi++) {
        float s = cq_scale(q, (size_t)row * groups + gi);
        size_t base = (size_t)row * q.pc + (size_t)gi * 128;
        for (int i = 0; i < 128; i++) g[i] = q.cb[cq_index(q.idx, base + i, q.bits)] * s;
        wht_inplace(g, 128);
        int n = std::min(128, q.cols - gi * 128);
        if (n > 0) std::memcpy(out + gi * 128, g, sizeof(float) * n);
    }
}

#if defined(__AVX2__)
#include <immintrin.h>

// Every AVX2 processor also has FMA3; MSVC enables both with /arch:AVX2 but defines no __FMA__.
#if defined(__FMA__) || defined(_MSC_VER)
#define OASIS_FMA 1
#endif

static inline __m256i dot_i8(__m256i acc, __m256i ax, __m256i ws) {
#if defined(__AVXVNNI__)
    return _mm256_dpbusd_avx_epi32(acc, ax, ws);
#else
    return _mm256_add_epi32(acc, _mm256_madd_epi16(_mm256_maddubs_epi16(ax, ws), _mm256_set1_epi16(1)));
#endif
}

static inline int hsum_i32(__m256i v) {
    __m128i s = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256(v, 1));
    s = _mm_add_epi32(s, _mm_shuffle_epi32(s, 0x4E));
    s = _mm_add_epi32(s, _mm_shuffle_epi32(s, 0xB1));
    return _mm_cvtsi128_si32(s);
}

#elif defined(OASIS_NEON)

// Signed int8 products of a and b added in fours into the int32 lanes of acc: SDOT where the core has the
// dot-product instructions, otherwise widening multiplies with pairwise accumulation. Both are exact.
static inline int32x4_t dot_i8(int32x4_t acc, int8x16_t a, int8x16_t b) {
#if defined(__ARM_FEATURE_DOTPROD)
    return vdotq_s32(acc, a, b);
#else
    acc = vpadalq_s16(acc, vmull_s8(vget_low_s8(a), vget_low_s8(b)));
    return vpadalq_s16(acc, vmull_high_s8(a, b));
#endif
}

#endif

bool g_exact_gemv = false;

void prep_x(const float* x, int cols, int bits, QX& o) {
    o.pc = (cols + 127) / 128 * 128;
    o.bits = bits;
    int groups = o.pc / 128;
    float xt[2048];
    std::memset(xt, 0, sizeof(float) * o.pc);
    std::memcpy(xt, x, sizeof(float) * cols);
    o.xs.resize(o.pc);
    o.xa.resize(o.pc);
    o.sx.resize(groups);
    for (int g = 0; g < groups; g++) {
        float* v = xt + g * 128;
        wht_inplace(v, 128);
#if defined(OASIS_NEON)
        int8_t* dst = o.xs.data() + g * 128;
        // The same scale, roundings and reordering as the loops below, sixteen values at a time.
        float32x4_t mx = vdupq_n_f32(0.0f);
        for (int i = 0; i < 128; i += 4) mx = vmaxnmq_f32(mx, vabsq_f32(vld1q_f32(v + i)));
        const float amax = vmaxnmvq_f32(mx);
        const float inv = amax > 0 ? 127.0f / amax : 0.0f;
        o.sx[g] = amax / 127.0f;
        int8_t tmp[128];
        for (int i = 0; i < 128; i += 16) {
            int32x4_t q[4];
            for (int k = 0; k < 4; k++) q[k] = vcvtnq_s32_f32(vmulq_n_f32(vld1q_f32(v + i + 4 * k), inv));
            const int16x8_t h0 = vcombine_s16(vmovn_s32(q[0]), vmovn_s32(q[1])), h1 = vcombine_s16(vmovn_s32(q[2]), vmovn_s32(q[3]));
            vst1q_s8(tmp + i, vcombine_s8(vmovn_s16(h0), vmovn_s16(h1)));
        }
        if (bits == 2) {
            const int8x16x4_t a = vld4q_s8(tmp), b = vld4q_s8(tmp + 64);
            for (int k = 0; k < 4; k++) {
                vst1q_s8(dst + k * 32, a.val[k]);
                vst1q_s8(dst + k * 32 + 16, b.val[k]);
            }
        } else {
            for (int c = 0; c < 2; c++) {
                const int8x16x2_t a = vld2q_s8(tmp + 64 * c), b = vld2q_s8(tmp + 64 * c + 32);
                vst1q_s8(dst + c * 64, a.val[0]);
                vst1q_s8(dst + c * 64 + 16, b.val[0]);
                vst1q_s8(dst + c * 64 + 32, a.val[1]);
                vst1q_s8(dst + c * 64 + 48, b.val[1]);
            }
        }
        for (int i = 0; i < 128; i += 16) vst1q_s8(o.xa.data() + g * 128 + i, vabsq_s8(vld1q_s8(dst + i)));
#else
        float amax = 0.0f;
        for (int i = 0; i < 128; i++) amax = std::max(amax, std::fabs(v[i]));
        float inv = amax > 0 ? 127.0f / amax : 0.0f;
        o.sx[g] = amax / 127.0f;
        int8_t tmp[128];
        for (int i = 0; i < 128; i++) tmp[i] = (int8_t)std::nearbyint(v[i] * inv);
        int8_t* dst = o.xs.data() + g * 128;
        if (bits == 2) {
            for (int k = 0; k < 4; k++)
                for (int mm = 0; mm < 32; mm++) dst[k * 32 + mm] = tmp[4 * mm + k];
        } else {
            for (int c = 0; c < 2; c++)
                for (int mm = 0; mm < 32; mm++) {
                    dst[c * 64 + mm] = tmp[2 * (32 * c + mm)];
                    dst[c * 64 + 32 + mm] = tmp[2 * (32 * c + mm) + 1];
                }
        }
        for (int i = 0; i < 128; i++) o.xa[g * 128 + i] = (int8_t)std::abs((int)dst[i]);
#endif
    }
}

CbTable cb_table(const CQ& q) {
    CbTable t;
    int ncb = 1 << q.bits;
    float cmax = 0.0f;
    for (int i = 0; i < ncb; i++) cmax = std::max(cmax, std::fabs(q.cb[i]));
    t.cs = cmax / 127.0f;
    std::memset(t.tbl, 0, sizeof t.tbl);
    for (int i = 0; i < ncb; i++) t.tbl[i] = (int8_t)std::nearbyint(q.cb[i] / t.cs);
    return t;
}

// One output of W x for a quantised row: int8 dot products per rotated group, scaled per group.
float cq_row(const CQ& q, const QX& qx, int r, const CbTable& t) {
    int groups = q.pc / 128;
    size_t row_bytes = (size_t)q.pc * q.bits / 8;
    const uint8_t* ib = q.idx + (size_t)r * row_bytes;
    float acc = 0.0f;
#if defined(__AVX2__)
    const __m256i T = _mm256_broadcastsi128_si256(_mm_load_si128((const __m128i*)t.tbl));
    const __m256i m3 = _mm256_set1_epi8(3), m15 = _mm256_set1_epi8(15);
    for (int g = 0; g < groups; g++) {
        const int8_t* xg = qx.xs.data() + g * 128;
        const int8_t* ag = qx.xa.data() + g * 128;
        __m256i a = _mm256_setzero_si256();
        if (q.bits == 2) {
            __m256i b = _mm256_loadu_si256((const __m256i*)(ib + g * 32));
            for (int k = 0; k < 4; k++) {
                __m256i w = _mm256_shuffle_epi8(T, _mm256_and_si256(_mm256_srli_epi16(b, 2 * k), m3));
                a = dot_i8(a, _mm256_loadu_si256((const __m256i*)(ag + k * 32)),
                           _mm256_sign_epi8(w, _mm256_loadu_si256((const __m256i*)(xg + k * 32))));
            }
        } else {
            for (int c = 0; c < 2; c++) {
                __m256i b = _mm256_loadu_si256((const __m256i*)(ib + g * 64 + c * 32));
                __m256i lo = _mm256_shuffle_epi8(T, _mm256_and_si256(b, m15));
                __m256i hi = _mm256_shuffle_epi8(T, _mm256_and_si256(_mm256_srli_epi16(b, 4), m15));
                a = dot_i8(a, _mm256_loadu_si256((const __m256i*)(ag + c * 64)),
                           _mm256_sign_epi8(lo, _mm256_loadu_si256((const __m256i*)(xg + c * 64))));
                a = dot_i8(a, _mm256_loadu_si256((const __m256i*)(ag + c * 64 + 32)),
                           _mm256_sign_epi8(hi, _mm256_loadu_si256((const __m256i*)(xg + c * 64 + 32))));
            }
        }
        acc += (float)hsum_i32(a) * qx.sx[g] * cq_scale(q, (size_t)r * groups + g);
    }
#elif defined(OASIS_NEON)
    // The same integer sums as the portable path below, sixteen weights per table lookup.
    const int8x16_t T = vld1q_s8(t.tbl);
    for (int g = 0; g < groups; g++) {
        const int8_t* xg = qx.xs.data() + g * 128;
        int32x4_t a = vdupq_n_s32(0);
        if (q.bits == 2) {
            const uint8x16_t m3 = vdupq_n_u8(3);
            const uint8x16_t b0 = vld1q_u8(ib + g * 32), b1 = vld1q_u8(ib + g * 32 + 16);
            a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(b0, m3)), vld1q_s8(xg));
            a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(b1, m3)), vld1q_s8(xg + 16));
            a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b0, 2), m3)), vld1q_s8(xg + 32));
            a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b1, 2), m3)), vld1q_s8(xg + 48));
            a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b0, 4), m3)), vld1q_s8(xg + 64));
            a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b1, 4), m3)), vld1q_s8(xg + 80));
            a = dot_i8(a, vqtbl1q_s8(T, vshrq_n_u8(b0, 6)), vld1q_s8(xg + 96));
            a = dot_i8(a, vqtbl1q_s8(T, vshrq_n_u8(b1, 6)), vld1q_s8(xg + 112));
        } else {
            const uint8x16_t m15 = vdupq_n_u8(15);
            for (int c = 0; c < 2; c++) {
                const uint8_t* bp = ib + g * 64 + c * 32;
                const int8_t* xp = xg + c * 64;
                const uint8x16_t b0 = vld1q_u8(bp), b1 = vld1q_u8(bp + 16);
                a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(b0, m15)), vld1q_s8(xp));
                a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(b1, m15)), vld1q_s8(xp + 16));
                a = dot_i8(a, vqtbl1q_s8(T, vshrq_n_u8(b0, 4)), vld1q_s8(xp + 32));
                a = dot_i8(a, vqtbl1q_s8(T, vshrq_n_u8(b1, 4)), vld1q_s8(xp + 48));
            }
        }
        acc += (float)vaddvq_s32(a) * qx.sx[g] * cq_scale(q, (size_t)r * groups + g);
    }
#else
    // Portable path: the same int8 arithmetic, one weight at a time.
    for (int g = 0; g < groups; g++) {
        const int8_t* xg = qx.xs.data() + g * 128;
        int a = 0;
        if (q.bits == 2) {
            for (int k = 0; k < 4; k++)
                for (int mm = 0; mm < 32; mm++) a += t.tbl[(ib[g * 32 + mm] >> (2 * k)) & 3] * xg[k * 32 + mm];
        } else {
            for (int c = 0; c < 2; c++)
                for (int mm = 0; mm < 32; mm++) {
                    uint8_t bb = ib[g * 64 + c * 32 + mm];
                    a += t.tbl[bb & 15] * xg[c * 64 + mm] + t.tbl[bb >> 4] * xg[c * 64 + 32 + mm];
                }
        }
        acc += (float)a * qx.sx[g] * cq_scale(q, (size_t)r * groups + g);
    }
#endif
    return acc * t.cs;
}

static void cq_rows_q(const CQ& q, const QX& qx, float* y, bool parallel) {
    CbTable t = cb_table(q);
    #pragma omp parallel for schedule(static) if (parallel)
    for (int r = 0; r < q.rows; r++) y[r] = q.zero_rows && q.zero_rows[r] ? 0.0f : cq_row(q, qx, r, t);
}

void cq_gemv_q(const CQ& q, const QX& qx, float* y) { cq_rows_q(q, qx, y, q.rows >= 256); }

// C[M, rows] = A[M, cols] W^T with int8 activations per rotated 128-group, as in the deployed kernels.
#if defined(OASIS_NEON)
// Every row of W against four prepared inputs at once: a group's weights are looked up once for all four,
// and each output gets the integer sums and the float steps of cq_row as GCC compiles it: with 8-bit scales
// the group's scale times the input's, then a multiply-add of that and the sum; with fp16 scales the sum
// times the input's scale, then a multiply-add of that and the group's scale.
static void cq_rows4(const CQ& q, const QX* qx, const CbTable& t, float* const* y) {
    const int groups = q.pc / 128;
    const size_t row_bytes = (size_t)q.pc * q.bits / 8;
    const int8x16_t T = vld1q_s8(t.tbl);
    const uint8x16_t m3 = vdupq_n_u8(3), m15 = vdupq_n_u8(15);
    for (int r = 0; r < q.rows; r++) {
        if (q.zero_rows && q.zero_rows[r]) {
            for (int i = 0; i < 4; i++) y[i][r] = 0.0f;
            continue;
        }
        const uint8_t* ib = q.idx + (size_t)r * row_bytes;
        float32x4_t acc = vdupq_n_f32(0.0f);
        for (int g = 0; g < groups; g++) {
            int8x16_t w[8];
            if (q.bits == 2) {
                const uint8x16_t b0 = vld1q_u8(ib + g * 32), b1 = vld1q_u8(ib + g * 32 + 16);
                w[0] = vqtbl1q_s8(T, vandq_u8(b0, m3));
                w[1] = vqtbl1q_s8(T, vandq_u8(b1, m3));
                w[2] = vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b0, 2), m3));
                w[3] = vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b1, 2), m3));
                w[4] = vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b0, 4), m3));
                w[5] = vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b1, 4), m3));
                w[6] = vqtbl1q_s8(T, vshrq_n_u8(b0, 6));
                w[7] = vqtbl1q_s8(T, vshrq_n_u8(b1, 6));
            } else {
                for (int c = 0; c < 2; c++) {
                    const uint8x16_t b0 = vld1q_u8(ib + g * 64 + c * 32), b1 = vld1q_u8(ib + g * 64 + c * 32 + 16);
                    w[4 * c] = vqtbl1q_s8(T, vandq_u8(b0, m15));
                    w[4 * c + 1] = vqtbl1q_s8(T, vandq_u8(b1, m15));
                    w[4 * c + 2] = vqtbl1q_s8(T, vshrq_n_u8(b0, 4));
                    w[4 * c + 3] = vqtbl1q_s8(T, vshrq_n_u8(b1, 4));
                }
            }
            int32x4_t a[4];
            float sx[4];
            for (int i = 0; i < 4; i++) {
                const int8_t* xg = qx[i].xs.data() + g * 128;
                int32x4_t s = vdupq_n_s32(0);
                for (int k = 0; k < 8; k++) s = dot_i8(s, w[k], vld1q_s8(xg + 16 * k));
                a[i] = s;
                sx[i] = qx[i].sx[g];
            }
            const float32x4_t sum = vcvtq_f32_s32(vpaddq_s32(vpaddq_s32(a[0], a[1]), vpaddq_s32(a[2], a[3])));
            const float scale = cq_scale(q, (size_t)r * groups + g);
            if (q.ncode) acc = vfmaq_f32(acc, vmulq_n_f32(vld1q_f32(sx), scale), sum);
            else acc = vfmaq_f32(acc, vmulq_f32(sum, vld1q_f32(sx)), vdupq_n_f32(scale));
        }
        float out[4];
        vst1q_f32(out, vmulq_n_f32(acc, t.cs));
        for (int i = 0; i < 4; i++) y[i][r] = out[i];
    }
}
#endif

void cq_gemm_q(const CQ& q, const float* A, int M, float* C) {
    CbTable t = cb_table(q);
#if defined(OASIS_NEON)
    const int blocks = (M + 3) / 4;
    #pragma omp parallel for schedule(static) if (M > 1)
    for (int b = 0; b < blocks; b++) {
        thread_local QX qx[4];
        const int i0 = b * 4, n = std::min(4, M - i0);
        for (int i = 0; i < n; i++) prep_x(A + (size_t)(i0 + i) * q.cols, q.cols, q.bits, qx[i]);
        if (n == 4) {
            float* y[4] = {C + (size_t)i0 * q.rows, C + (size_t)(i0 + 1) * q.rows, C + (size_t)(i0 + 2) * q.rows, C + (size_t)(i0 + 3) * q.rows};
            cq_rows4(q, qx, t, y);
            continue;
        }
        for (int i = 0; i < n; i++) {
            float* y = C + (size_t)(i0 + i) * q.rows;
            for (int r = 0; r < q.rows; r++) y[r] = q.zero_rows && q.zero_rows[r] ? 0.0f : cq_row(q, qx[i], r, t);
        }
    }
#else
    #pragma omp parallel for schedule(static) if (M > 1)
    for (int i = 0; i < M; i++) {
        thread_local QX qx;
        prep_x(A + (size_t)i * q.cols, q.cols, q.bits, qx);
        float* y = C + (size_t)i * q.rows;
        for (int r = 0; r < q.rows; r++) y[r] = q.zero_rows && q.zero_rows[r] ? 0.0f : cq_row(q, qx, r, t);
    }
#endif
}

float dotf(const float* a, const float* b, int n) {
#if defined(OASIS_FMA)
    __m256 s0 = _mm256_setzero_ps(), s1 = _mm256_setzero_ps();
    int i = 0;
    for (; i + 16 <= n; i += 16) {
        s0 = _mm256_fmadd_ps(_mm256_loadu_ps(a + i), _mm256_loadu_ps(b + i), s0);
        s1 = _mm256_fmadd_ps(_mm256_loadu_ps(a + i + 8), _mm256_loadu_ps(b + i + 8), s1);
    }
    s0 = _mm256_add_ps(s0, s1);
    __m128 h = _mm_add_ps(_mm256_castps256_ps128(s0), _mm256_extractf128_ps(s0, 1));
    h = _mm_add_ps(h, _mm_movehl_ps(h, h));
    h = _mm_add_ss(h, _mm_shuffle_ps(h, h, 1));
    float s = _mm_cvtss_f32(h);
    for (; i < n; i++) s += a[i] * b[i];
    return s;
#elif defined(OASIS_NEON)
    float32x4_t s0 = vdupq_n_f32(0.0f), s1 = s0, s2 = s0, s3 = s0;
    int i = 0;
    for (; i + 16 <= n; i += 16) {
        s0 = vfmaq_f32(s0, vld1q_f32(a + i), vld1q_f32(b + i));
        s1 = vfmaq_f32(s1, vld1q_f32(a + i + 4), vld1q_f32(b + i + 4));
        s2 = vfmaq_f32(s2, vld1q_f32(a + i + 8), vld1q_f32(b + i + 8));
        s3 = vfmaq_f32(s3, vld1q_f32(a + i + 12), vld1q_f32(b + i + 12));
    }
    float s = vaddvq_f32(vaddq_f32(vaddq_f32(s0, s1), vaddq_f32(s2, s3)));
    for (; i < n; i++) s += a[i] * b[i];
    return s;
#else
    float s = 0.0f;
    for (int i = 0; i < n; i++) s += a[i] * b[i];
    return s;
#endif
}

void cq_gemv(const CQ& q, const float* x, float* y) {
#if defined(__AVX2__) || defined(OASIS_NEON)
    if (!g_exact_gemv && (q.bits == 2 || q.bits == 4)) {
        thread_local QX qx;
        prep_x(x, q.cols, q.bits, qx);
        cq_gemv_q(q, qx, y);
        return;
    }
#endif
    int groups = q.pc / 128;
    std::vector<float> xt(q.pc, 0.0f);
    std::memcpy(xt.data(), x, sizeof(float) * q.cols);
    for (int gi = 0; gi < groups; gi++) wht_inplace(xt.data() + gi * 128, 128);
    if (q.bits == 2) {
        alignas(16) float tab[256][4];
        for (int b = 0; b < 256; b++)
            for (int k = 0; k < 4; k++) tab[b][k] = q.cb[(b >> (2 * k)) & 3];
        size_t row_bytes = (size_t)q.pc / 4;
        #pragma omp parallel for schedule(static) if (q.rows >= 256)
        for (int r = 0; r < q.rows; r++) {
            const uint8_t* ib = q.idx + (size_t)r * row_bytes;
            float acc = 0.0f;
            for (int gi = 0; gi < groups; gi++) {
                const uint8_t* gb = ib + gi * 32;
                const float* xg = xt.data() + gi * 128;
                float a0 = 0, a1 = 0, a2 = 0, a3 = 0;
                for (int m = 0; m < 32; m++) {
                    const float* t = tab[gb[m]];
                    const float* xv = xg + 4 * m;
                    a0 += t[0] * xv[0];
                    a1 += t[1] * xv[1];
                    a2 += t[2] * xv[2];
                    a3 += t[3] * xv[3];
                }
                acc += (a0 + a1 + a2 + a3) * cq_scale(q, (size_t)r * groups + gi);
            }
            y[r] = acc;
        }
        return;
    }
    if (q.bits == 4) {
        alignas(8) float tab[256][2];
        for (int b = 0; b < 256; b++) {
            tab[b][0] = q.cb[b & 15];
            tab[b][1] = q.cb[b >> 4];
        }
        size_t row_bytes = (size_t)q.pc / 2;
        #pragma omp parallel for schedule(static) if (q.rows >= 256)
        for (int r = 0; r < q.rows; r++) {
            const uint8_t* ib = q.idx + (size_t)r * row_bytes;
            float acc = 0.0f;
            for (int gi = 0; gi < groups; gi++) {
                const uint8_t* gb = ib + gi * 64;
                const float* xg = xt.data() + gi * 128;
                float a0 = 0, a1 = 0;
                for (int m = 0; m < 64; m++) {
                    a0 += tab[gb[m]][0] * xg[2 * m];
                    a1 += tab[gb[m]][1] * xg[2 * m + 1];
                }
                acc += (a0 + a1) * cq_scale(q, (size_t)r * groups + gi);
            }
            y[r] = acc;
        }
        return;
    }
    #pragma omp parallel for schedule(static) if (q.rows >= 256)
    for (int r = 0; r < q.rows; r++) {
        float acc = 0.0f;
        for (int gi = 0; gi < groups; gi++) {
            float a = 0.0f;
            size_t base = (size_t)r * q.pc + (size_t)gi * 128;
            for (int i = 0; i < 128; i++) a += q.cb[cq_index(q.idx, base + i, q.bits)] * xt[gi * 128 + i];
            acc += a * cq_scale(q, (size_t)r * groups + gi);
        }
        y[r] = acc;
    }
}

void gemm_nt(const float* A, int M, int K, const float* W, int N, float* C) {
    if (M == 1) {
        #pragma omp parallel for schedule(static) if ((long)N * K > 200000)
        for (int j = 0; j < N; j++) C[j] = dotf(A, W + (size_t)j * K, K);
        return;
    }
    #pragma omp parallel for schedule(static) if ((long)M * N * K > 200000)
    for (int i = 0; i < M; i++)
        for (int j = 0; j < N; j++) C[(size_t)i * N + j] = dotf(A + (size_t)i * K, W + (size_t)j * K, K);
}

// C = A B. Every output accumulates its K products in order (fused multiply-adds where the target
// has them), so the register-blocked path and the plain loop give the same values.
static void gemm_nn_rows(const float* A, int i0, int i1, int K, const float* B, int N, float* C, int j0) {
    for (int i = i0; i < i1; i++) {
        float* c = C + (size_t)i * N;
        const float* a = A + (size_t)i * K;
        for (int j = j0; j < N; j++) c[j] = 0.0f;
        for (int k = 0; k < K; k++) {
            float av = a[k];
            const float* b = B + (size_t)k * N;
            for (int j = j0; j < N; j++) c[j] += av * b[j];
        }
    }
}

void gemm_nn(const float* A, int M, int K, const float* B, int N, float* C) {
#if defined(OASIS_FMA)
    const int NB = N / 16 * 16;
    const int blocks = (M + 3) / 4;
    #pragma omp parallel for schedule(static) if ((long)M * N * K > 200000)
    for (int ib = 0; ib < blocks; ib++) {
        int i0 = ib * 4, mr = std::min(4, M - i0);
        if (mr == 4) {
            const float* a0 = A + (size_t)i0 * K;
            const float* a1 = a0 + K;
            const float* a2 = a1 + K;
            const float* a3 = a2 + K;
            for (int j0 = 0; j0 < NB; j0 += 16) {
                __m256 c00 = _mm256_setzero_ps(), c01 = _mm256_setzero_ps(), c10 = _mm256_setzero_ps(), c11 = _mm256_setzero_ps();
                __m256 c20 = _mm256_setzero_ps(), c21 = _mm256_setzero_ps(), c30 = _mm256_setzero_ps(), c31 = _mm256_setzero_ps();
                for (int k = 0; k < K; k++) {
                    const float* b = B + (size_t)k * N + j0;
                    __m256 b0 = _mm256_loadu_ps(b), b1 = _mm256_loadu_ps(b + 8), x;
                    x = _mm256_broadcast_ss(a0 + k);
                    c00 = _mm256_fmadd_ps(x, b0, c00);
                    c01 = _mm256_fmadd_ps(x, b1, c01);
                    x = _mm256_broadcast_ss(a1 + k);
                    c10 = _mm256_fmadd_ps(x, b0, c10);
                    c11 = _mm256_fmadd_ps(x, b1, c11);
                    x = _mm256_broadcast_ss(a2 + k);
                    c20 = _mm256_fmadd_ps(x, b0, c20);
                    c21 = _mm256_fmadd_ps(x, b1, c21);
                    x = _mm256_broadcast_ss(a3 + k);
                    c30 = _mm256_fmadd_ps(x, b0, c30);
                    c31 = _mm256_fmadd_ps(x, b1, c31);
                }
                float* c = C + (size_t)i0 * N + j0;
                _mm256_storeu_ps(c, c00);
                _mm256_storeu_ps(c + 8, c01);
                _mm256_storeu_ps(c + N, c10);
                _mm256_storeu_ps(c + N + 8, c11);
                _mm256_storeu_ps(c + 2 * N, c20);
                _mm256_storeu_ps(c + 2 * N + 8, c21);
                _mm256_storeu_ps(c + 3 * N, c30);
                _mm256_storeu_ps(c + 3 * N + 8, c31);
            }
        } else {
            for (int i = i0; i < i0 + mr; i++) {
                const float* a = A + (size_t)i * K;
                for (int j0 = 0; j0 < NB; j0 += 16) {
                    __m256 c0 = _mm256_setzero_ps(), c1 = _mm256_setzero_ps();
                    for (int k = 0; k < K; k++) {
                        const float* b = B + (size_t)k * N + j0;
                        __m256 x = _mm256_broadcast_ss(a + k);
                        c0 = _mm256_fmadd_ps(x, _mm256_loadu_ps(b), c0);
                        c1 = _mm256_fmadd_ps(x, _mm256_loadu_ps(b + 8), c1);
                    }
                    _mm256_storeu_ps(C + (size_t)i * N + j0, c0);
                    _mm256_storeu_ps(C + (size_t)i * N + j0 + 8, c1);
                }
            }
        }
        if (NB < N) gemm_nn_rows(A, i0, i0 + mr, K, B, N, C, NB);
    }
#elif defined(OASIS_NEON)
    const int NB = N / 16 * 16;
    const int blocks = (M + 3) / 4;
    #pragma omp parallel for schedule(static) if ((long)M * N * K > 200000)
    for (int ib = 0; ib < blocks; ib++) {
        int i0 = ib * 4, mr = std::min(4, M - i0);
        if (mr == 4) {
            const float* a0 = A + (size_t)i0 * K;
            const float* a1 = a0 + K;
            const float* a2 = a1 + K;
            const float* a3 = a2 + K;
            for (int j0 = 0; j0 < NB; j0 += 16) {
                float32x4_t c00 = vdupq_n_f32(0.0f), c01 = c00, c02 = c00, c03 = c00, c10 = c00, c11 = c00, c12 = c00, c13 = c00;
                float32x4_t c20 = c00, c21 = c00, c22 = c00, c23 = c00, c30 = c00, c31 = c00, c32 = c00, c33 = c00;
                for (int k = 0; k < K; k++) {
                    const float* b = B + (size_t)k * N + j0;
                    const float32x4_t b0 = vld1q_f32(b), b1 = vld1q_f32(b + 4), b2 = vld1q_f32(b + 8), b3 = vld1q_f32(b + 12);
                    float x = a0[k];
                    c00 = vfmaq_n_f32(c00, b0, x); c01 = vfmaq_n_f32(c01, b1, x); c02 = vfmaq_n_f32(c02, b2, x); c03 = vfmaq_n_f32(c03, b3, x);
                    x = a1[k];
                    c10 = vfmaq_n_f32(c10, b0, x); c11 = vfmaq_n_f32(c11, b1, x); c12 = vfmaq_n_f32(c12, b2, x); c13 = vfmaq_n_f32(c13, b3, x);
                    x = a2[k];
                    c20 = vfmaq_n_f32(c20, b0, x); c21 = vfmaq_n_f32(c21, b1, x); c22 = vfmaq_n_f32(c22, b2, x); c23 = vfmaq_n_f32(c23, b3, x);
                    x = a3[k];
                    c30 = vfmaq_n_f32(c30, b0, x); c31 = vfmaq_n_f32(c31, b1, x); c32 = vfmaq_n_f32(c32, b2, x); c33 = vfmaq_n_f32(c33, b3, x);
                }
                float* c = C + (size_t)i0 * N + j0;
                vst1q_f32(c, c00); vst1q_f32(c + 4, c01); vst1q_f32(c + 8, c02); vst1q_f32(c + 12, c03);
                c += N;
                vst1q_f32(c, c10); vst1q_f32(c + 4, c11); vst1q_f32(c + 8, c12); vst1q_f32(c + 12, c13);
                c += N;
                vst1q_f32(c, c20); vst1q_f32(c + 4, c21); vst1q_f32(c + 8, c22); vst1q_f32(c + 12, c23);
                c += N;
                vst1q_f32(c, c30); vst1q_f32(c + 4, c31); vst1q_f32(c + 8, c32); vst1q_f32(c + 12, c33);
            }
        } else {
            for (int i = i0; i < i0 + mr; i++) {
                const float* a = A + (size_t)i * K;
                for (int j0 = 0; j0 < NB; j0 += 16) {
                    float32x4_t c0 = vdupq_n_f32(0.0f), c1 = c0, c2 = c0, c3 = c0;
                    for (int k = 0; k < K; k++) {
                        const float* b = B + (size_t)k * N + j0;
                        const float x = a[k];
                        c0 = vfmaq_n_f32(c0, vld1q_f32(b), x);
                        c1 = vfmaq_n_f32(c1, vld1q_f32(b + 4), x);
                        c2 = vfmaq_n_f32(c2, vld1q_f32(b + 8), x);
                        c3 = vfmaq_n_f32(c3, vld1q_f32(b + 12), x);
                    }
                    float* c = C + (size_t)i * N + j0;
                    vst1q_f32(c, c0); vst1q_f32(c + 4, c1); vst1q_f32(c + 8, c2); vst1q_f32(c + 12, c3);
                }
            }
        }
        if (NB < N) gemm_nn_rows(A, i0, i0 + mr, K, B, N, C, NB);
    }
#else
    #pragma omp parallel for schedule(static) if ((long)M * N * K > 200000)
    for (int i = 0; i < M; i++) gemm_nn_rows(A, i, i + 1, K, B, N, C, 0);
#endif
}

void zc_norm(const float* x, int n, const float* scale, float* out) {
    double ss = 0.0;
    for (int i = 0; i < n; i++) ss += (double)x[i] * x[i];
    float r = 1.0f / std::sqrt((float)(ss / n) + 1e-6f);
    for (int i = 0; i < n; i++) out[i] = x[i] * r * (1.0f + scale[i]);
}

void rms_unit(const float* x, int n, float* out) {
    double ss = 0.0;
    for (int i = 0; i < n; i++) ss += (double)x[i] * x[i];
    float r = 1.0f / std::sqrt((float)(ss / n) + 1e-6f);
    for (int i = 0; i < n; i++) out[i] = x[i] * r;
}

void band65(const float* u, const float* w, int n, float* out) {
    for (int j = 0; j < n; j++) {
        const float* wj = w + (size_t)j * 65;
        int lo = std::max(0, 32 - j), hi = std::min(65, n + 32 - j);
        float s = 0.0f;
        for (int t = lo; t < hi; t++) s += wj[t] * u[j + t - 32];
        out[j] = s;
    }
}

void qmm(Model& m, const std::string& name, const float* A, int M, float* C) {
    CQ q = m.Q(name);
    if (g_exact_gemv) gemm_nt(A, M, q.cols, m.Wdq(name), q.rows, C);
    else cq_gemm_q(q, A, M, C);
}

}  // namespace oasis

// ==== Image side: Pillow-exact bilinear resize, encoder, neck, detection, masks, box slots and memory

namespace oasis {

// ------------------------------------------------------------------ resize (matches Pillow 8bpc)
// The resize follows the fixed-point bilinear filter of Pillow, so that the engine reads the same
// pixels as the reference implementation. Pillow carries this notice:
//
//   The Python Imaging Library (PIL) is
//
//       Copyright © 1997-2011 by Secret Labs AB
//       Copyright © 1995-2011 by Fredrik Lundh and contributors
//
//   Pillow is the friendly PIL fork. It is
//
//       Copyright © 2010 by Jeffrey A. Clark and contributors
//
//   Like PIL, Pillow is licensed under the open source MIT-CMU License:
//
//   By obtaining, using, and/or copying this software and/or its associated
//   documentation, you agree that you have read, understood, and will comply
//   with the following terms and conditions:
//
//   Permission to use, copy, modify and distribute this software and its
//   documentation for any purpose and without fee is hereby granted,
//   provided that the above copyright notice appears in all copies, and that
//   both that copyright notice and this permission notice appear in supporting
//   documentation, and that the name of Secret Labs AB or the author not be
//   used in advertising or publicity pertaining to distribution of the software
//   without specific, written prior permission.
//
//   SECRET LABS AB AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS
//   SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS.
//   IN NO EVENT SHALL SECRET LABS AB OR THE AUTHOR BE LIABLE FOR ANY SPECIAL,
//   INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
//   LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
//   OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
//   PERFORMANCE OF THIS SOFTWARE.
static const int PREC = 22;

static void coeffs(int in_size, int out_size, std::vector<int>& bounds, std::vector<int>& kk, int& ksize) {
    double scale = (double)in_size / out_size, fs = scale < 1.0 ? 1.0 : scale;
    double support = 1.0 * fs;
    ksize = (int)std::ceil(support) * 2 + 1;
    bounds.assign(out_size * 2, 0);
    kk.assign((size_t)out_size * ksize, 0);
    std::vector<double> k(ksize);
    for (int xx = 0; xx < out_size; xx++) {
        double center = (xx + 0.5) * scale, ww = 0.0, ss = 1.0 / fs;
        int xmin = (int)(center - support + 0.5);
        if (xmin < 0) xmin = 0;
        int xmax = (int)(center + support + 0.5);
        if (xmax > in_size) xmax = in_size;
        xmax -= xmin;
        for (int x = 0; x < xmax; x++) {
            double t = (x + xmin - center + 0.5) * ss;
            double w = t < 0 ? -t : t;
            w = w < 1.0 ? 1.0 - w : 0.0;
            k[x] = w;
            ww += w;
        }
        for (int x = 0; x < xmax; x++) {
            double v = ww != 0.0 ? k[x] / ww : 0.0;
            kk[(size_t)xx * ksize + x] = v < 0 ? (int)(-0.5 + v * (1 << PREC)) : (int)(0.5 + v * (1 << PREC));
        }
        bounds[xx * 2] = xmin;
        bounds[xx * 2 + 1] = xmax;
    }
}

static inline uint8_t clip8(long long v) {
    if (v >= ((long long)256 << PREC)) return 255;
    if (v <= 0) return 0;
    return (uint8_t)(v >> PREC);
}

std::vector<uint8_t> resize_rgb(const uint8_t* src, int W0, int H0, int W1, int H1) {
    std::vector<uint8_t> tmp((size_t)H0 * W1 * 3), out((size_t)H1 * W1 * 3);
    std::vector<int> b, k;
    int ks;
    if (W1 != W0) {
        coeffs(W0, W1, b, k, ks);
        #pragma omp parallel for schedule(static) if ((long)H0 * W1 > 65536)
        for (int y = 0; y < H0; y++)
            for (int x = 0; x < W1; x++)
                for (int c = 0; c < 3; c++) {
                    long long ss = 1LL << (PREC - 1);
                    for (int t = 0; t < b[x * 2 + 1]; t++) ss += (long long)src[((size_t)y * W0 + b[x * 2] + t) * 3 + c] * k[(size_t)x * ks + t];
                    tmp[((size_t)y * W1 + x) * 3 + c] = clip8(ss);
                }
    } else {
        std::memcpy(tmp.data(), src, tmp.size());
    }
    if (H1 != H0) {
        coeffs(H0, H1, b, k, ks);
        #pragma omp parallel for schedule(static) if ((long)H1 * W1 > 65536)
        for (int y = 0; y < H1; y++)
            for (int x = 0; x < W1; x++)
                for (int c = 0; c < 3; c++) {
                    long long ss = 1LL << (PREC - 1);
                    for (int t = 0; t < b[y * 2 + 1]; t++) ss += (long long)tmp[((size_t)(b[y * 2] + t) * W1 + x) * 3 + c] * k[(size_t)y * ks + t];
                    out[((size_t)y * W1 + x) * 3 + c] = clip8(ss);
                }
    } else {
        out = tmp;
    }
    return out;
}

// ------------------------------------------------------------------ NHWC building blocks
static std::vector<uint8_t> valid_mask(int n, int vh, int vw) {
    std::vector<uint8_t> m((size_t)n * n, 0);
    for (int i = 0; i < std::min(n, vh); i++)
        for (int j = 0; j < std::min(n, vw); j++) m[(size_t)i * n + j] = 1;
    return m;
}

// 3x3 convolution with one pixel of zero padding on every side, HWIO weights.
static void conv3x3(const float* in, int H, int W, int Ci, const float* w, int Co, int sh, int sw, std::vector<float>& out, int& Ho, int& Wo) {
    Ho = (H - 1) / sh + 1;
    Wo = (W - 1) / sw + 1;
    out.assign((size_t)Ho * Wo * Co, 0.0f);
    #pragma omp parallel for schedule(static)
    for (int y = 0; y < Ho; y++)
        for (int x = 0; x < Wo; x++) {
            float* o = out.data() + ((size_t)y * Wo + x) * Co;
            for (int kh = 0; kh < 3; kh++) {
                int iy = y * sh - 1 + kh;
                if (iy < 0 || iy >= H) continue;
                for (int kw = 0; kw < 3; kw++) {
                    int ix = x * sw - 1 + kw;
                    if (ix < 0 || ix >= W) continue;
                    const float* ip = in + ((size_t)iy * W + ix) * Ci;
                    const float* wp = w + ((size_t)(kh * 3 + kw) * Ci) * Co;
                    for (int ci = 0; ci < Ci; ci++) {
                        float v = ip[ci];
                        const float* wr = wp + (size_t)ci * Co;
                        for (int co = 0; co < Co; co++) o[co] += v * wr[co];
                    }
                }
            }
        }
}

// Depthwise kh x kw convolution, weights [kh, kw, 1, C]; padding kh/2, kw/2.
static void dwconv(const float* in, int H, int W, int C, const float* w, int kh_, int kw_, int sh, int sw, std::vector<float>& out, int& Ho, int& Wo) {
    int ph = kh_ / 2, pw_ = kw_ / 2;
    Ho = (H + 2 * ph - kh_) / sh + 1;
    Wo = (W + 2 * pw_ - kw_) / sw + 1;
    out.assign((size_t)Ho * Wo * C, 0.0f);
    #pragma omp parallel for schedule(static)
    for (int y = 0; y < Ho; y++)
        for (int x = 0; x < Wo; x++) {
            float* o = out.data() + ((size_t)y * Wo + x) * C;
            for (int a = 0; a < kh_; a++) {
                int iy = y * sh - ph + a;
                if (iy < 0 || iy >= H) continue;
                for (int bb = 0; bb < kw_; bb++) {
                    int ix = x * sw - pw_ + bb;
                    if (ix < 0 || ix >= W) continue;
                    const float* ip = in + ((size_t)iy * W + ix) * C;
                    const float* wp = w + (size_t)(a * kw_ + bb) * C;
                    for (int c = 0; c < C; c++) o[c] += ip[c] * wp[c];
                }
            }
        }
}

// Fixed 4096-element blocks, so every element is computed the same way whatever the thread count.
static void silu_v(std::vector<float>& v) {
    const int n = (int)v.size(), B = 4096, nb = (n + B - 1) / B;
    float* p = v.data();
    #pragma omp parallel for schedule(static) if (n > 65536)
    for (int b = 0; b < nb; b++) {
        int e = std::min(n, (b + 1) * B);
#if defined(OASIS_NEON)
        const int i0 = b * B, n4 = (e - i0) & ~3;
        silu_inplace(p + i0, n4);
        for (int i = i0 + n4; i < e; i++) p[i] = siluf(p[i]);
#else
        for (int i = b * B; i < e; i++) p[i] = siluf(p[i]);
#endif
    }
}

static void norm_rows(std::vector<float>& x, int C, const float* scale) {
    size_t n = x.size() / C;
    #pragma omp parallel for schedule(static) if (n > 64)
    for (long p = 0; p < (long)n; p++) zc_norm(x.data() + p * C, C, scale, x.data() + p * C);
}

static void apply_mask(std::vector<float>& x, int C, const std::vector<uint8_t>& m) {
    for (size_t p = 0; p < m.size(); p++)
        if (!m[p]) std::memset(x.data() + p * C, 0, sizeof(float) * C);
}

static std::vector<float> pointwise(const std::vector<float>& x, int Ci, const float* k, int Co) {
    int P = (int)(x.size() / Ci);
    std::vector<float> out((size_t)P * Co);
    gemm_nn(x.data(), P, Ci, k, Co, out.data());
    return out;
}

static std::vector<float> up2(const std::vector<float>& x, int n, int C) {
    std::vector<float> o((size_t)4 * n * n * C);
    for (int i = 0; i < 2 * n; i++)
        for (int j = 0; j < 2 * n; j++)
            std::memcpy(o.data() + ((size_t)i * 2 * n + j) * C, x.data() + ((size_t)(i / 2) * n + j / 2) * C, sizeof(float) * C);
    return o;
}

// One HConvGLU's weights, looked up once per block rather than once per cell.
struct HConv {
    const float *gate, *gate_bias, *value, *out;
};

static HConv hconv_of(Model& m, const std::string& pre) {
    return {m.F(pre + "/gate_conv"), m.F(pre + "/gate_bias"), m.F(pre + "/value_conv"), m.F(pre + "/out_conv")};
}

#if defined(OASIS_NEON)
// The bands of NC cells at once, with tap-major weights w [65, 512] (Model::load stores them so on Arm):
// every tap is one multiply-add across the outputs of all NC cells, so each weight vector is loaded once
// for all of them. up[c] holds a cell's input with 32 zeros on either side.
template <int NC>
static void band_cells(float (*up)[576], const float* w, float (*out)[512]) {
    for (int j = 0; j < 512; j += 8) {
        float32x4_t a[NC][2];
        for (int c = 0; c < NC; c++) a[c][0] = a[c][1] = vdupq_n_f32(0.0f);
        for (int t = 0; t < 65; t++) {
            const float32x4_t w0 = vld1q_f32(w + (size_t)t * 512 + j), w1 = vld1q_f32(w + (size_t)t * 512 + j + 4);
            for (int c = 0; c < NC; c++) {
                a[c][0] = vfmaq_f32(a[c][0], w0, vld1q_f32(up[c] + j + t));
                a[c][1] = vfmaq_f32(a[c][1], w1, vld1q_f32(up[c] + j + t + 4));
            }
        }
        for (int c = 0; c < NC; c++) {
            vst1q_f32(out[c] + j, a[c][0]);
            vst1q_f32(out[c] + j + 4, a[c][1]);
        }
    }
}

static void band_cells(int nc, float (*up)[576], const float* w, float (*out)[512]) {
    switch (nc) {
        case 4: band_cells<4>(up, w, out); break;
        case 3: band_cells<3>(up, w, out); break;
        case 2: band_cells<2>(up, w, out); break;
        default: band_cells<1>(up, w, out); break;
    }
}

// HConvGLU of up to four cells at once: x[c] -> y[c].
static void hconvglu_cells(const Model& m, const HConv& w, const float* const* x, float* const* y, int nc) {
    float up[4][576], a[4][512], v[4][512], z[4][512];
    for (int c = 0; c < nc; c++) {
        std::memset(up[c], 0, sizeof up[c]);
        wht_seq512(m, x[c], up[c] + 32);
    }
    band_cells(nc, up, w.gate, a);
    band_cells(nc, up, w.value, v);
    // gelu_tanh(a + bias) * v with the operations the scalar loop compiles to, k * fma(0.044715 x, x^2, x) and
    // ((1 + tanh) v)(0.5 x), and tanh four at a time.
    for (int c = 0; c < nc; c++)
        for (int i = 0; i < 512; i += 4) {
            const float32x4_t x = vaddq_f32(vld1q_f32(a[c] + i), vld1q_f32(w.gate_bias + i));
            const float32x4_t arg = vmulq_n_f32(vfmaq_f32(x, vmulq_n_f32(x, 0.044715f), vmulq_f32(x, x)), 0.7978845608028654f);
            const float32x4_t t = vaddq_f32(tanhf4(arg), vdupq_n_f32(1.0f));
            vst1q_f32(up[c] + 32 + i, vmulq_f32(vmulq_f32(t, vld1q_f32(v[c] + i)), vmulq_n_f32(x, 0.5f)));
        }
    band_cells(nc, up, w.out, z);
    for (int c = 0; c < nc; c++) wht_seq512(m, z[c], y[c]);
}
#else
static void hconvglu(const Model& m, const HConv& w, const float* x, float* y) {
    float u[512], a[512], v[512], h[512], z[512];
    wht_seq512(m, x, u);
    band65(u, w.gate, 512, a);
    band65(u, w.value, 512, v);
    for (int i = 0; i < 512; i++) h[i] = gelu_tanh(a[i] + w.gate_bias[i]) * v[i];
    band65(h, w.out, 512, z);
    wht_seq512(m, z, y);
}
#endif

// residual block B(F) = F + PW(DW3x3(SiLU(N(F)))), masked
static void res_block(Model& m, std::vector<float>& F, int n, int C, const std::string& norm, const std::string& dw,
                      const std::string& pw, const std::vector<uint8_t>& mask) {
    std::vector<float> t = F, d;
    norm_rows(t, C, m.F(norm));
    silu_v(t);
    int ho, wo;
    dwconv(t.data(), n, n, C, m.F(dw), 3, 3, 1, 1, d, ho, wo);
    std::vector<float> p = pointwise(d, C, m.F(pw), C);
    for (size_t i = 0; i < F.size(); i++) F[i] += p[i];
    apply_mask(F, C, mask);
}

// Stage checksums on stderr when OASIS_DEBUG=1, for comparing builds.
static void stage(const char* name, const std::vector<float>& v) {
    static const bool on = [] { const char* d = std::getenv("OASIS_DEBUG"); return d && d[0] == '1'; }();
    if (!on) return;
    double s = 0.0, a = 0.0;
    for (float x : v) { s += x; a += std::fabs(x); }
    std::fprintf(stderr, "%-10s n=%zu sum=%.6g abs=%.6g first=%.6g\n", name, v.size(), s, a, v.empty() ? 0.0 : v[0]);
}

Encoded encode(Model& m, const uint8_t* rgb, int W0, int H0) {
    Encoded e;
    float r = 512.0f / (float)std::max(W0, H0);
    int w = std::max(1, (int)std::nearbyint((double)W0 * r)), h = std::max(1, (int)std::nearbyint((double)H0 * r));
    e.meta = {W0, H0, w, h, r};
    std::vector<uint8_t> rs = (w == W0 && h == H0) ? std::vector<uint8_t>(rgb, rgb + (size_t)W0 * H0 * 3) : resize_rgb(rgb, W0, H0, w, h);
    std::vector<float> canvas((size_t)512 * 512 * 3, 0.0f);
    #pragma omp parallel for schedule(static)
    for (int y = 0; y < h; y++)
        for (int x = 0; x < w; x++)
            for (int c = 0; c < 3; c++) canvas[((size_t)y * 512 + x) * 3 + c] = rs[((size_t)y * w + x) * 3 + c] / 127.5f - 1.0f;
    auto vm = [&](int s) { return valid_mask(512 / s, (h + s - 1) / s, (w + s - 1) / s); };
    auto m2 = vm(2), m4 = vm(4);
    e.valid8 = vm(8);
    e.valid16 = vm(16);
    e.valid32 = vm(32);
    e.valid4 = m4;

    stage("canvas", canvas);
    std::vector<float> f, t, d;
    int Ho, Wo;
    conv3x3(canvas.data(), 512, 512, 3, m.F("stem/w0"), 32, 2, 2, f, Ho, Wo);
    norm_rows(f, 32, m.F("stem/norm_0/scale"));
    apply_mask(f, 32, m2);
    int C = 32, n = 256;
    const int chans[5] = {32, 64, 128, 256, 512};
    std::vector<float> feats[6];
    for (int j = 1; j <= 4; j++) {
        t = f;
        silu_v(t);
        dwconv(t.data(), n, n, C, m.F("stem/dw_" + std::to_string(j)), 3, 3, 2, 2, d, Ho, Wo);
        f = pointwise(d, C, m.F("stem/pw_" + std::to_string(j) + "/kernel"), chans[j]);
        C = chans[j];
        n = Ho;
        norm_rows(f, C, m.F("stem/norm_" + std::to_string(j) + "/scale"));
        apply_mask(f, C, j == 1 ? m4 : j == 2 ? e.valid8 : j == 3 ? e.valid16 : e.valid32);
        feats[j] = f;
    }
    for (int j = 1; j <= 4; j++) stage(("stem" + std::to_string(j)).c_str(), feats[j]);
    e.S4 = feats[1];
    silu_v(e.S4);
    std::vector<float> S8 = feats[2], S16 = feats[3], X = feats[4];
    silu_v(S8);
    silu_v(S16);
    silu_v(X);

    // ---- convolutional body, 16 x 16 x 512
    const auto& v32 = e.valid32;
    for (int i = 0; i < 16; i++) {
        std::string pre = "body_" + std::to_string(i);
        std::vector<float> Z = X, va, ga, mix;
        norm_rows(Z, 512, m.F(pre + "/mix_norm/scale"));
        dwconv(Z.data(), 16, 16, 512, m.F(pre + "/mix/dw_value"), 3, 3, 1, 1, va, Ho, Wo);
        dwconv(Z.data(), 16, 16, 512, m.F(pre + "/mix/dw_gate"), 3, 3, 1, 1, ga, Ho, Wo);
        const float* bv = m.F(pre + "/mix/b_value");
        const float* bg = m.F(pre + "/mix/b_gate");
        #pragma omp parallel for schedule(static)
        for (int p = 0; p < 256; p++)
            for (int c = 0; c < 512; c++) {
                size_t k = (size_t)p * 512 + c;
                va[k] = (va[k] + bv[c]) * siluf(ga[k] + bg[c]);
            }
        dwconv(va.data(), 16, 16, 512, m.F(pre + "/mix/dw_down"), 3, 3, 1, 1, mix, Ho, Wo);
        for (int p = 0; p < 256; p++)
            if (v32[p])
                for (int c = 0; c < 512; c++) X[(size_t)p * 512 + c] += mix[(size_t)p * 512 + c];
        std::vector<float> Hh((size_t)256 * 512);
        const float* hs = m.F(pre + "/hconv_norm/scale");
        const HConv hw = hconv_of(m, pre + "/hconv");
#if defined(OASIS_NEON)
        #pragma omp parallel for schedule(static)
        for (int b = 0; b < 64; b++) {
            float xn[4][512];
            const float* xs[4];
            float* ys[4];
            for (int c = 0; c < 4; c++) {
                const int p = b * 4 + c;
                zc_norm(X.data() + (size_t)p * 512, 512, hs, xn[c]);
                xs[c] = xn[c];
                ys[c] = Hh.data() + (size_t)p * 512;
            }
            hconvglu_cells(m, hw, xs, ys, 4);
        }
#else
        #pragma omp parallel for schedule(static)
        for (int p = 0; p < 256; p++) {
            float xn[512];
            zc_norm(X.data() + (size_t)p * 512, 512, hs, xn);
            hconvglu(m, hw, xn, Hh.data() + (size_t)p * 512);
        }
#endif
        const float* gam = m.F(pre + "/grn/gamma");
        const float* bet = m.F(pre + "/grn/beta");
        std::vector<float> rc(512);
        #pragma omp parallel for schedule(static)
        for (int c = 0; c < 512; c++) {
            double s = 0.0;
            for (int p = 0; p < 256; p++)
                if (v32[p]) s += (double)Hh[(size_t)p * 512 + c] * Hh[(size_t)p * 512 + c];
            rc[c] = (float)std::sqrt(s + 1e-6);
        }
        double rbar = 0.0;
        for (int c = 0; c < 512; c++) rbar += rc[c];
        float rb = (float)(rbar / 512.0);
        #pragma omp parallel for schedule(static)
        for (int p = 0; p < 256; p++)
            if (v32[p])
                for (int c = 0; c < 512; c++) {
                    float hv = Hh[(size_t)p * 512 + c];
                    X[(size_t)p * 512 + c] += hv + gam[c] * hv * rc[c] / (rb + 1e-6f) + bet[c];
                }
    }

    stage("body", X);
    // ---- attention blocks over 256 cells + 4 registers
    const int T = 260;
    std::vector<float> Z((size_t)T * 512);
    std::memcpy(Z.data(), X.data(), sizeof(float) * 256 * 512);
    std::memcpy(Z.data() + 256 * 512, m.F("registers"), sizeof(float) * 4 * 512);
    std::vector<uint8_t> tv(T, 1);
    for (int p = 0; p < 256; p++) tv[p] = v32[p];
#if defined(OASIS_NEON)
    std::vector<int> active;   // the valid tokens, in order, for the four-at-a-time HConvGLU
    for (int p = 0; p < T; p++)
        if (tv[p]) active.push_back(p);
#endif
    std::vector<float> E((size_t)256 * 512, 0.0f);
    auto add_level = [&](const std::vector<float>& Zl, int l) {
        const float* sc = m.F("level_norm_" + std::to_string(l) + "/scale");
        const float* g = m.F("level_gates") + (size_t)l * 512;
        #pragma omp parallel for schedule(static)
        for (int p = 0; p < 256; p++) {
            float tmp[512];
            zc_norm(Zl.data() + (size_t)p * 512, 512, sc, tmp);
            for (int c = 0; c < 512; c++) E[(size_t)p * 512 + c] += g[c] * tmp[c];
        }
    };
    float cosr[256][16], sinr[256][16];
    for (int p = 0; p < 256; p++)
        for (int r2 = 0; r2 < 8; r2++) {
            double fr = 1.0 / std::pow(10000.0, (2.0 * r2) / 16.0);
            double ai = (p / 16) * fr, aj = (p % 16) * fr;
            cosr[p][r2] = (float)std::cos(ai);
            sinr[p][r2] = (float)std::sin(ai);
            cosr[p][8 + r2] = (float)std::cos(aj);
            sinr[p][8 + r2] = (float)std::sin(aj);
        }
    std::vector<float> z((size_t)T * 512), q((size_t)T * 128), k((size_t)T * 32), v((size_t)T * 64), gt((size_t)T * 256), o((size_t)T * 256), A((size_t)T * 512);
    for (int l = 0; l < 12; l++) {
        add_level(Z, l);
        std::string pre = "top_" + std::to_string(l) + "/";
        const float* an = m.F(pre + "attn_norm/scale");
        for (int p = 0; p < T; p++) zc_norm(Z.data() + (size_t)p * 512, 512, an, z.data() + (size_t)p * 512);
        qmm(m, pre + "self_attn/q_proj/kernel", z.data(), T, q.data());
        qmm(m, pre + "self_attn/k_proj/kernel", z.data(), T, k.data());
        qmm(m, pre + "self_attn/v_proj/kernel", z.data(), T, v.data());
        qmm(m, pre + "self_attn/gate_proj/kernel", z.data(), T, gt.data());
        const float* qn = m.F(pre + "self_attn/q_norm/scale");
        const float* kn = m.F(pre + "self_attn/k_norm/scale");
        auto rope = [&](float* x, int p) {
            if (p >= 256) return;
            float x1[16], x2[16];
            std::memcpy(x1, x, sizeof x1);
            std::memcpy(x2, x + 16, sizeof x2);
            for (int i = 0; i < 16; i++) {
                x[i] = x1[i] * cosr[p][i] - x2[i] * sinr[p][i];
                x[16 + i] = x2[i] * cosr[p][i] + x1[i] * sinr[p][i];
            }
        };
        for (int p = 0; p < T; p++) {
            for (int hh = 0; hh < 4; hh++) {
                float* qp = q.data() + (size_t)p * 128 + hh * 32;
                zc_norm(qp, 32, qn, qp);
                rope(qp, p);
            }
            float* kp = k.data() + (size_t)p * 32;
            zc_norm(kp, 32, kn, kp);
            rope(kp, p);
        }
        #pragma omp parallel for schedule(static)
        for (int p = 0; p < T; p++) {
#if defined(OASIS_NEON)
            std::vector<float> s(T), e(T);
#else
            std::vector<float> s(T);
#endif
            for (int hh = 0; hh < 4; hh++) {
                const float* qp = q.data() + (size_t)p * 128 + hh * 32;
                float mx = -1e30f;
                for (int j = 0; j < T; j++) {
                    if (!tv[j]) { s[j] = -1e30f; continue; }
                    float d = 0.0f;
                    for (int c = 0; c < 32; c++) d += qp[c] * k[(size_t)j * 32 + c];
                    s[j] = d / std::sqrt(32.0f);
                    mx = std::max(mx, s[j]);
                }
                double den = 0.0;
#if defined(OASIS_NEON)
                for (int j = 0; j < T; j += 4) vst1q_f32(e.data() + j, expf4(vsubq_f32(vld1q_f32(s.data() + j), vdupq_n_f32(mx))));
                for (int j = 0; j < T; j++) {
                    s[j] = tv[j] ? e[j] : 0.0f;
                    den += s[j];
                }
#else
                for (int j = 0; j < T; j++) {
                    s[j] = tv[j] ? std::exp(s[j] - mx) : 0.0f;
                    den += s[j];
                }
#endif
                float* op = o.data() + (size_t)p * 256 + hh * 64;
                for (int c = 0; c < 64; c++) op[c] = 0.0f;
                for (int j = 0; j < T; j++) {
                    if (s[j] == 0.0f) continue;
                    float wv = (float)(s[j] / den);
                    for (int c = 0; c < 64; c++) op[c] += wv * v[(size_t)j * 64 + c];
                }
#if defined(OASIS_NEON)
                mul_sigmoid(op, gt.data() + (size_t)p * 256 + hh * 64, 64);
#else
                for (int c = 0; c < 64; c++) op[c] *= sigmoidf(gt[(size_t)p * 256 + hh * 64 + c]);
#endif
            }
        }
        qmm(m, pre + "self_attn/out_proj/kernel", o.data(), T, A.data());
        float ag = sigmoidf(m.F(pre + "attn_gate")[0]);
        const float* pn = m.F(pre + "post_attn_norm/scale");
        const float* hn = m.F(pre + "hconv_norm/scale");
        const HConv hw = hconv_of(m, pre + "hconv");
#if defined(OASIS_NEON)
        const int nb = ((int)active.size() + 3) / 4;
        #pragma omp parallel for schedule(static)
        for (int b = 0; b < nb; b++) {
            const int nc = std::min(4, (int)active.size() - b * 4);
            float tmp[4][512], hy[4][512];
            const float* xs[4];
            float* ys[4];
            for (int c = 0; c < nc; c++) {
                float* zp = Z.data() + (size_t)active[b * 4 + c] * 512;
                float t1[512];
                zc_norm(A.data() + (size_t)active[b * 4 + c] * 512, 512, pn, t1);
                for (int k = 0; k < 512; k++) zp[k] += ag * t1[k];
                zc_norm(zp, 512, hn, tmp[c]);
                xs[c] = tmp[c];
                ys[c] = hy[c];
            }
            hconvglu_cells(m, hw, xs, ys, nc);
            for (int c = 0; c < nc; c++) {
                float* zp = Z.data() + (size_t)active[b * 4 + c] * 512;
                for (int k = 0; k < 512; k++) zp[k] += hy[c][k];
            }
        }
#else
        #pragma omp parallel for schedule(static)
        for (int p = 0; p < T; p++) {
            if (!tv[p]) continue;
            float tmp[512], hy[512];
            float* zp = Z.data() + (size_t)p * 512;
            zc_norm(A.data() + (size_t)p * 512, 512, pn, tmp);
            for (int c = 0; c < 512; c++) zp[c] += ag * tmp[c];
            zc_norm(zp, 512, hn, tmp);
            hconvglu(m, hw, tmp, hy);
            for (int c = 0; c < 512; c++) zp[c] += hy[c];
        }
#endif
    }
    {
        const float* en = m.F("enc_norm/scale");
        float tmp[512];
        for (int p = 0; p < 256; p++) {
            zc_norm(Z.data() + (size_t)p * 512, 512, en, tmp);
            for (int c = 0; c < 512; c++) E[(size_t)p * 512 + c] += tmp[c];
        }
        apply_mask(E, 512, v32);
    }
    stage("E", E);
    e.E = E;

    // ---- neck to D [64, 64, 128]
    std::vector<float> co = pointwise(E, 512, m.F("dense/coarse/kernel"), 128);
    std::vector<float> tt = up2(co, 16, 128), mid = pointwise(S16, 256, m.F("dense/mid/kernel"), 128);
    for (size_t i = 0; i < tt.size(); i++) tt[i] += mid[i];
    apply_mask(tt, 128, e.valid16);
    res_block(m, tt, 32, 128, "dense/neck_norm/scale", "dense/neck_dw", "dense/neck_pw/kernel", e.valid16);
    std::vector<float> dd = up2(tt, 32, 128), fine = pointwise(S8, 128, m.F("dense/fine/kernel"), 128);
    for (size_t i = 0; i < dd.size(); i++) dd[i] += fine[i];
    apply_mask(dd, 128, e.valid8);
    for (int i = 0; i < 3; i++)
        res_block(m, dd, 64, 128, "dense/block_norm_" + std::to_string(i) + "/scale", "dense/dw_" + std::to_string(i),
                  "dense/pw_" + std::to_string(i) + "/kernel", e.valid8);
    norm_rows(dd, 128, m.F("dense/final_norm/scale"));
    silu_v(dd);
    apply_mask(dd, 128, e.valid8);
    stage("D", dd);
    e.D = dd;
    e.det = pointwise(dd, 128, m.F("dense/out/kernel"), 10);
    const float* db = m.F("dense/bias");
    for (int p = 0; p < 64 * 64; p++)
        for (int c = 0; c < 10; c++) e.det[(size_t)p * 10 + c] += db[c];
    stage("det", e.det);
    return e;
}

// ------------------------------------------------------------------ detection
static float iou_box(const Box& a, const Box& b) {
    float ix = std::max(0.0f, std::min(a.x1, b.x1) - std::max(a.x0, b.x0));
    float iy = std::max(0.0f, std::min(a.y1, b.y1) - std::max(a.y0, b.y0));
    float inter = ix * iy;
    float u = (a.x1 - a.x0) * (a.y1 - a.y0) + (b.x1 - b.x0) * (b.y1 - b.y0) - inter;
    return u > 0 ? inter / u : 0.0f;
}

void detect(const Encoded& enc, float conf, float text_conf, bool nms, int max_obj, int max_txt,
            std::vector<Box>& objects, std::vector<Box>& texts) {
    for (int ch = 0; ch < 2; ch++) {
        float thr = ch == 0 ? conf : text_conf;
        int limit = ch == 0 ? max_obj : max_txt;
        std::vector<float> heat(64 * 64);
        for (int p = 0; p < 64 * 64; p++) heat[p] = sigmoidf(enc.det[(size_t)p * 10 + ch]);
        struct Cand { float s; int i, j; };
        std::vector<Cand> c;
        for (int i = 0; i < 64; i++)
            for (int j = 0; j < 64; j++) {
                int p = i * 64 + j;
                if (!enc.valid8[p] || heat[p] < thr) continue;
                bool peak = true;
                for (int di = -1; di <= 1 && peak; di++)
                    for (int dj = -1; dj <= 1; dj++) {
                        int a = i + di, b = j + dj;
                        if (a < 0 || b < 0 || a >= 64 || b >= 64) continue;
                        if (heat[a * 64 + b] > heat[p]) { peak = false; break; }
                    }
                if (peak) c.push_back({heat[p], i, j});
            }
        std::stable_sort(c.begin(), c.end(), [](const Cand& a, const Cand& b) { return a.s > b.s; });
        std::vector<Box>& out = ch == 0 ? objects : texts;
        float vw = enc.meta.w / 512.0f, vh = enc.meta.h / 512.0f;
        for (auto& cd : c) {
            const float* d = enc.det.data() + ((size_t)cd.i * 64 + cd.j) * 10;
            float w = std::fabs(d[2 + 2 * ch]), h = std::fabs(d[3 + 2 * ch]);
            float cx = cd.j + d[6 + 2 * ch], cy = cd.i + d[7 + 2 * ch];
            Box b{(cx - w / 2) / 64, (cy - h / 2) / 64, (cx + w / 2) / 64, (cy + h / 2) / 64, cd.s, cd.i, cd.j};
            b.x0 = std::clamp(b.x0, 0.0f, vw);
            b.x1 = std::clamp(b.x1, 0.0f, vw);
            b.y0 = std::clamp(b.y0, 0.0f, vh);
            b.y1 = std::clamp(b.y1, 0.0f, vh);
            bool dup = false;
            if (nms)
                for (auto& x : out)
                    if (iou_box(b, x) > 0.5f) { dup = true; break; }
            if (dup) continue;
            out.push_back(b);
            if ((int)out.size() >= limit) break;
        }
    }
}

// ------------------------------------------------------------------ masks
static std::vector<float> head_branch(Model& m, std::vector<float> F, int n, const std::string& pre,
                                      const std::vector<uint8_t>& mask, const std::string& outk) {
    for (int i = 0; i < 2; i++)
        res_block(m, F, n, 128, "dense/" + pre + "_norm_" + std::to_string(i) + "/scale", "dense/" + pre + "_dw_" + std::to_string(i),
                  "dense/" + pre + "_pw_" + std::to_string(i) + "/kernel", mask);
    norm_rows(F, 128, m.F("dense/" + pre + "_final_norm/scale"));
    silu_v(F);
    return pointwise(F, 128, m.F(outk), 32);
}

// Mask coefficients at one centre cell. The coefficient branch's only spatial operations are two 3x3
// depthwise convolutions, so its output at a cell depends on the 5 x 5 neighbourhood of D alone; cells
// outside the grid enter as masked zeros, as the full branch's padding does.
static std::vector<float> coef_at(Model& m, const Encoded& enc, int ci, int cj) {
    std::vector<float> F((size_t)25 * 128, 0.0f);
    std::vector<uint8_t> pm(25, 0);
    for (int di = -2; di <= 2; di++)
        for (int dj = -2; dj <= 2; dj++) {
            int i = ci + di, j = cj + dj, p = (di + 2) * 5 + (dj + 2);
            if (i < 0 || j < 0 || i >= 64 || j >= 64) continue;
            pm[p] = enc.valid8[(size_t)i * 64 + j];
            std::memcpy(F.data() + (size_t)p * 128, enc.D.data() + ((size_t)i * 64 + j) * 128, sizeof(float) * 128);
        }
    std::vector<float> C = head_branch(m, F, 5, "coef", pm, "dense/coef/kernel");
    return std::vector<float>(C.begin() + 12 * 32, C.begin() + 13 * 32);
}

std::vector<float> mask_prob(Model& m, const Encoded& enc, const Box& b, std::vector<float>* protos) {
    if (protos->empty()) {
        std::vector<float> U = up2(enc.D, 64, 128), ed = pointwise(enc.S4, 64, m.F("dense/edge/kernel"), 128);
        for (size_t i = 0; i < U.size(); i++) U[i] += ed[i];
        apply_mask(U, 128, enc.valid4);
        *protos = head_branch(m, U, 128, "proto", enc.valid4, "dense/protos/kernel");
    }
    const float* P = protos->data();
    std::vector<float> coef = coef_at(m, enc, b.ci, b.cj);
    const float* a = coef.data();
    std::vector<float> out(128 * 128, 0.0f);
    int ci = std::min(127, (int)((b.y0 + b.y1) / 2 * 128)), cj = std::min(127, (int)((b.x0 + b.x1) / 2 * 128));
    for (int i = 0; i < 128; i++) {
        float yc = (i + 0.5f) / 128;
        for (int j = 0; j < 128; j++) {
            float xc = (j + 0.5f) / 128;
            bool inside = (yc >= b.y0 && yc <= b.y1 && xc >= b.x0 && xc <= b.x1) || (i == ci && j == cj);
            if (!inside) continue;
            const float* pp = P + ((size_t)i * 128 + j) * 32;
            float s = 0.0f;
            for (int r = 0; r < 32; r++) s += a[r] * pp[r];
            out[(size_t)i * 128 + j] = sigmoidf(s);
        }
    }
    return out;
}

// ------------------------------------------------------------------ box slots and decoder memory
void memory(Model& m, const Encoded& enc, const std::vector<Box>& objects, const std::vector<Box>& texts,
            std::vector<float>& M, std::vector<uint8_t>& mvalid, std::vector<int>& refs) {
    float pe = sigmoidf(m.F("pe_gate")[0]);
    M.assign((size_t)256 * 512, 0.0f);
    for (int p = 0; p < 256; p++) {
        int ri = p / 16, cj = p % 16;
        for (int half = 0; half < 2; half++) {
            int pos = half == 0 ? ri : cj;
            for (int i = 0; i < 128; i++) {
                double om = 1.0 / std::pow(10000.0, (double)i / 128.0);
                M[(size_t)p * 512 + half * 256 + i] = (float)std::sin(pos * om);
                M[(size_t)p * 512 + half * 256 + 128 + i] = (float)std::cos(pos * om);
            }
        }
        for (int c = 0; c < 512; c++) M[(size_t)p * 512 + c] = enc.E[(size_t)p * 512 + c] + pe * M[(size_t)p * 512 + c];
    }
    mvalid.assign(enc.valid32.begin(), enc.valid32.end());
    refs.clear();
    const float* sp = m.F("slot_proj/kernel");
    const float* sn = m.F("slot_norm/scale");
    const float* emb_scale_dummy = nullptr;
    (void)emb_scale_dummy;
    CQ emb = m.Q("embedding/embedding");
    std::vector<float> row(512);
    for (int kind = 0; kind < 2; kind++) {
        const auto& boxes = kind == 0 ? objects : texts;
        for (size_t k = 0; k < boxes.size(); k++) {
            const Box& b = boxes[k];
            std::vector<float> feat(192);
            std::memcpy(feat.data(), enc.D.data() + ((size_t)b.ci * 64 + b.cj) * 128, sizeof(float) * 128);
            float bc[4] = {b.x0, b.y0, b.x1, b.y1};
            for (int c = 0; c < 4; c++)
                for (int r2 = 0; r2 < 8; r2++) {
                    double a = std::pow(2.0, r2) * 3.14159265358979323846 * bc[c];
                    feat[128 + c * 16 + r2] = (float)std::sin(a);
                    feat[128 + c * 16 + 8 + r2] = (float)std::cos(a);
                }
            std::vector<float> pooled(512, 0.0f);
            int count = 0, ci = std::min(15, (int)((b.y0 + b.y1) / 2 * 16)), cc = std::min(15, (int)((b.x0 + b.x1) / 2 * 16));
            std::vector<uint8_t> in(256, 0);
            for (int i = 0; i < 16; i++)
                for (int j = 0; j < 16; j++) {
                    float yc = (i + 0.5f) / 16, xc = (j + 0.5f) / 16;
                    in[i * 16 + j] = yc >= b.y0 && yc <= b.y1 && xc >= b.x0 && xc <= b.x1;
                }
            in[ci * 16 + cc] = 1;
            for (int p = 0; p < 256; p++) in[p] = in[p] && enc.valid32[p];
            bool any = false;
            for (int p = 0; p < 256; p++) any = any || in[p];
            if (!any) in[ci * 16 + cc] = 1;
            for (int p = 0; p < 256; p++)
                if (in[p]) {
                    count++;
                    for (int c = 0; c < 512; c++) pooled[c] += enc.E[(size_t)p * 512 + c];
                }
            for (int c = 0; c < 512; c++) pooled[c] /= count;
            std::vector<float> z(512, 0.0f), zp(512);
            gemm_nn(feat.data(), 1, 192, sp, 512, z.data());
            qmm(m, "slot_pool/kernel", pooled.data(), 1, zp.data());
            for (int c = 0; c < 512; c++) z[c] += zp[c];
            std::vector<float> s(512);
            zc_norm(z.data(), 512, sn, s.data());
            int rid = m.ref_base + (kind == 0 ? (int)k : 64 + (int)k);
            cq_dequant_row(emb, rid, row.data());
            for (int c = 0; c < 512; c++) s[c] += std::sqrt(512.0f) * row[c];
            M.insert(M.end(), s.begin(), s.end());
            mvalid.push_back(1);
            refs.push_back(rid);
        }
    }
}

}  // namespace oasis

// ==== Text side: CTC reader, decoder with caches, detokenisation

namespace oasis {

// ------------------------------------------------------------------ reader
static void conv3x3_s(const float* in, int H, int W, int Ci, const float* w, int Co, int sh, int sw, std::vector<float>& out, int& Ho, int& Wo) {
    Ho = (H - 1) / sh + 1;
    Wo = (W - 1) / sw + 1;
    out.assign((size_t)Ho * Wo * Co, 0.0f);
    for (int y = 0; y < Ho; y++)
        for (int x = 0; x < Wo; x++) {
            float* o = out.data() + ((size_t)y * Wo + x) * Co;
            for (int kh = 0; kh < 3; kh++) {
                int iy = y * sh - 1 + kh;
                if (iy < 0 || iy >= H) continue;
                for (int kw = 0; kw < 3; kw++) {
                    int ix = x * sw - 1 + kw;
                    if (ix < 0 || ix >= W) continue;
                    const float* ip = in + ((size_t)iy * W + ix) * Ci;
                    const float* wp = w + ((size_t)(kh * 3 + kw) * Ci) * Co;
                    for (int ci = 0; ci < Ci; ci++)
                        for (int co = 0; co < Co; co++) o[co] += ip[ci] * wp[(size_t)ci * Co + co];
                }
            }
        }
}

static void dw_s(const float* in, int H, int W, int C, const float* w, int kh_, int kw_, int sh, int sw, std::vector<float>& out, int& Ho, int& Wo) {
    int ph = kh_ / 2, pw_ = kw_ / 2;
    Ho = (H + 2 * ph - kh_) / sh + 1;
    Wo = (W + 2 * pw_ - kw_) / sw + 1;
    out.assign((size_t)Ho * Wo * C, 0.0f);
    for (int y = 0; y < Ho; y++)
        for (int x = 0; x < Wo; x++) {
            float* o = out.data() + ((size_t)y * Wo + x) * C;
            for (int a = 0; a < kh_; a++) {
                int iy = y * sh - ph + a;
                if (iy < 0 || iy >= H) continue;
                for (int b = 0; b < kw_; b++) {
                    int ix = x * sw - pw_ + b;
                    if (ix < 0 || ix >= W) continue;
                    const float* ip = in + ((size_t)iy * W + ix) * C;
                    const float* wp = w + (size_t)(a * kw_ + b) * C;
                    for (int c = 0; c < C; c++) o[c] += ip[c] * wp[c];
                }
            }
        }
}

static void norm_all(std::vector<float>& x, int C, const float* s) {
    for (size_t p = 0; p < x.size() / C; p++) zc_norm(x.data() + p * C, C, s, x.data() + p * C);
}

// One 48 x 320 window -> argmax label per frame (80 frames).
static std::vector<int> reader_window(Model& m, const std::vector<float>& J) {
    std::vector<float> G, t, d;
    int H = 48, W = 320, Ho, Wo;
    conv3x3_s(J.data(), H, W, 3, m.F("reader/w0"), 32, 2, 2, G, Ho, Wo);
    H = Ho;
    W = Wo;
    norm_all(G, 32, m.F("reader/stage_norm_0/scale"));
    const int Cs[4] = {32, 64, 128, 384};
    const int sh[3] = {2, 2, 2}, sw[3] = {2, 1, 1};
    for (int j = 0; j < 3; j++) {
        t = G;
        for (auto& v : t) v = siluf(v);
        dw_s(t.data(), H, W, Cs[j], m.F("reader/dw_" + std::to_string(j)), 3, 3, sh[j], sw[j], d, Ho, Wo);
        H = Ho;
        W = Wo;
        G.assign((size_t)H * W * Cs[j + 1], 0.0f);
        qmm(m, "reader/pw_" + std::to_string(j) + "/kernel", d.data(), H * W, G.data());
        norm_all(G, Cs[j + 1], m.F("reader/stage_norm_" + std::to_string(j + 1) + "/scale"));
    }
    // H == 3, W == 80, C == 384
    std::vector<float> Z((size_t)W * 384, 0.0f);
    for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++)
            for (int c = 0; c < 384; c++) Z[(size_t)x * 384 + c] += siluf(G[((size_t)y * W + x) * 384 + c]) / H;
    std::vector<float> h, mixed, vg((size_t)W * 768), gated((size_t)W * 384), dn((size_t)W * 384);
    for (int b = 0; b < 4; b++) {
        h = Z;
        norm_all(h, 384, m.F("reader/norm_" + std::to_string(b) + "/scale"));
        int ho, wo;
        dw_s(h.data(), 1, W, 384, m.F("reader/mix_" + std::to_string(b)), 1, 5, 1, 1, mixed, ho, wo);
        qmm(m, "reader/up_" + std::to_string(b) + "/kernel", mixed.data(), W, vg.data());
        for (int x = 0; x < W; x++)
            for (int c = 0; c < 384; c++) gated[(size_t)x * 384 + c] = vg[(size_t)x * 768 + c] * siluf(vg[(size_t)x * 768 + 384 + c]);
        qmm(m, "reader/down_" + std::to_string(b) + "/kernel", gated.data(), W, dn.data());
        for (size_t i = 0; i < Z.size(); i++) Z[i] += dn[i];
    }
    norm_all(Z, 384, m.F("reader/final_norm/scale"));
    std::vector<float> lg((size_t)W * 307);
    qmm(m, "reader/out/kernel", Z.data(), W, lg.data());
    const float* bias = m.F("reader/out/bias");
    std::vector<int> lab(W);
    for (int x = 0; x < W; x++) {
        int best = 0;
        float bv = -1e30f;
        for (int c = 0; c < 307; c++) {
            float v = lg[(size_t)x * 307 + c] + bias[c];
            if (v > bv) { bv = v; best = c; }
        }
        lab[x] = best;
    }
    return lab;
}

static void utf8_append(std::string& s, uint32_t cp) {
    if (cp < 0x80) s += (char)cp;
    else if (cp < 0x800) { s += (char)(0xC0 | (cp >> 6)); s += (char)(0x80 | (cp & 0x3F)); }
    else if (cp < 0x10000) { s += (char)(0xE0 | (cp >> 12)); s += (char)(0x80 | ((cp >> 6) & 0x3F)); s += (char)(0x80 | (cp & 0x3F)); }
    else { s += (char)(0xF0 | (cp >> 18)); s += (char)(0x80 | ((cp >> 12) & 0x3F)); s += (char)(0x80 | ((cp >> 6) & 0x3F)); s += (char)(0x80 | (cp & 0x3F)); }
}

std::string read_text(Model& m, const uint8_t* rgb, int W, int H, float x0, float y0, float x1, float y1) {
    int cx0 = (int)std::floor(x0), cy0 = (int)std::floor(y0), cx1 = (int)std::ceil(x1), cy1 = (int)std::ceil(y1);
    int cw = cx1 - cx0, ch = cy1 - cy0;
    if (cw < 1 || ch < 1) return "";
    std::vector<uint8_t> crop((size_t)cw * ch * 3, 0);
    for (int y = 0; y < ch; y++)
        for (int x = 0; x < cw; x++) {
            int sy = cy0 + y, sx = cx0 + x;
            if (sy < 0 || sx < 0 || sy >= H || sx >= W) continue;
            std::memcpy(&crop[((size_t)y * cw + x) * 3], &rgb[((size_t)sy * W + sx) * 3], 3);
        }
    int U = std::max(1, (int)std::nearbyint(48.0 * cw / ch));
    std::vector<uint8_t> strip = resize_rgb(crop.data(), cw, ch, U, 48);
    std::vector<int> starts;
    if (U <= 320) starts.push_back(0);
    else for (int s = 0; s < U - 320 + 223; s += 224) starts.push_back(s);
    int n_frames = (U + 3) / 4;
    std::string out;
    int prev = 0;
    for (int xs : starts) {
        std::vector<float> J((size_t)48 * 320 * 3, 127.0f / 127.5f - 1.0f);
        for (int y = 0; y < 48; y++)
            for (int x = 0; x < 320 && xs + x < U; x++)
                for (int c = 0; c < 3; c++) J[((size_t)y * 320 + x) * 3 + c] = strip[((size_t)y * U + xs + x) * 3 + c] / 127.5f - 1.0f;
        std::vector<int> lab = reader_window(m, J);
        int lo = xs == 0 ? 0 : 12;
        int hi = xs + 320 >= U ? (U - xs + 3) / 4 : 68;
        for (int t = lo; t < std::min(hi, 80); t++) {
            if (xs / 4 + t >= n_frames) break;
            int l = lab[t];
            if (l != 0 && l != prev) utf8_append(out, m.charset[l - 1]);
            prev = l;
        }
    }
    return out;
}

// ------------------------------------------------------------------ decoder
static const int ORDERS[2] = {2, 3};
static const uint32_t SEED = 0x9E3779B9u, PRIME = 0x01000193u;

static uint32_t engram_index(const std::vector<int>& toks, int t, int oi, int order, int h) {
    uint32_t acc = SEED * (uint32_t)(oi * 2 + h + 1);
    for (int j = 0; j < order; j++) {
        uint32_t tok = t - j >= 0 ? (uint32_t)toks[t - j] : 0u;
        acc = (acc ^ tok) * PRIME;
    }
    acc ^= acc >> 15;
    return acc % 18432u;
}

static inline float fexp(float x) {
    x = x < -87.0f ? -87.0f : x;
    float t = x * 1.4426950408889634f, fi = std::floor(t), f = t - fi;
    float p = 1.0f + f * (0.69314718f + f * (0.24022651f + f * (0.05550411f + f * (0.00961813f + f * (0.00133336f + f * 0.00015404f)))));
    int32_t bits = ((int32_t)fi + 127) << 23;
    float sc;
    std::memcpy(&sc, &bits, 4);
    return p * sc;
}

// 4 x 4 Sinkhorn: twenty rounds of row then column normalisation, multiplicative form.
static void sinkhorn(float* L) {
    double K[16];
    float mx = L[0];
    for (int i = 1; i < 16; i++) mx = std::max(mx, L[i]);
    for (int i = 0; i < 16; i++) K[i] = std::exp((double)L[i] - mx);
    for (int it = 0; it < 20; it++) {
        for (int i = 0; i < 4; i++) {
            double s = K[i * 4] + K[i * 4 + 1] + K[i * 4 + 2] + K[i * 4 + 3];
            for (int j = 0; j < 4; j++) K[i * 4 + j] /= s;
        }
        for (int j = 0; j < 4; j++) {
            double s = K[j] + K[4 + j] + K[8 + j] + K[12 + j];
            for (int i = 0; i < 4; i++) K[i * 4 + j] /= s;
        }
    }
    for (int i = 0; i < 16; i++) L[i] = (float)K[i];
}

#if defined(OASIS_NEON)
static inline uint16_t to_half(float f) {
    __fp16 h = (__fp16)f;
    uint16_t b;
    std::memcpy(&b, &h, 2);
    return b;
}
#endif

struct Decoder {
    Model& m;
    int nM;
    std::vector<uint8_t> mvalid;
    // Image-memory keys and values, shared by copies. On Arm, whose caches are small, they are kept in fp16,
    // which halves what each token reads, and laid out in the order a head reads them: keys per head as
    // [nMp / 16][48][16] (memory entries padded to nMp), values per head as [nM][64].
    struct Cross {
        std::vector<float> k[8], v[8];
        std::vector<uint16_t> kt[8], v16[8];
    };
    std::shared_ptr<Cross> cross = std::make_shared<Cross>();
    int nMp = 0;
    std::vector<float>* ck = cross->k;
    std::vector<float>* cv = cross->v;
    std::vector<int> toks;
    // Self-attention keys and values of every position so far: [t][96] and [t][128], both groups interleaved;
    // on Arm one array per group, [t][48] and [t][64], so that a head reads its group's history in one run.
#if defined(OASIS_NEON)
    std::vector<float> K[8][2], V[8][2];
    std::vector<float> raw[8][3];
#else
    std::vector<float> K[8], V[8], raw[8][3];
#endif
    std::vector<float> ev[2];
    CQ emb;
    struct Layer {
        CQ q, k, v, gate, out, cq_, cgate, cout;
        const float *taps[3], *qn, *kn, *pre, *post, *cross, *pcross, *cqn, *hada;
        float ag, cg;
        const float *d1, *d2, *d3, *d4, *b2, *cond_v, *cond_u, *w1a, *w1b, *w2a, *w2b, *w3a, *w3b;
    } Ls[8];
    struct Site { CQ tab, key, val; const float* taps; } sites[2];
    const float *phi_pre, *phi_post, *phi_res, *a_pre, *a_post, *a_res, *b_pre, *b_post, *b_res, *final_norm;
    QX qx;
    SpinPool& pool;

    std::string L(const char* name, int l) { return std::string("stack/layers/block/") + name + "/" + std::to_string(l); }

    // y_i = W_i x for several quantised matrices of one width, x already prepared in qx; rows are shared out.
    void rows(std::initializer_list<const CQ*> mats, std::initializer_list<float*> outs, int chunk = 32) {
        const CQ* q[4];
        float* y[4];
        int off[5] = {0}, n = 0;
        for (const CQ* p : mats) q[n++] = p;
        n = 0;
        for (float* p : outs) y[n++] = p;
        for (int i = 0; i < n; i++) off[i + 1] = off[i] + q[i]->rows;
#if defined(OASIS_NEON)
        // An Arm core reads a long run of rows far faster than scattered pieces of a few kilobytes, so each
        // thread claims one contiguous share of the rows; the result does not depend on how rows are shared.
        chunk = std::max(chunk, (off[n] + pool.size() - 1) / pool.size());
#endif
        const CbTable t = cb_table(*q[0]);
        pool.run(off[n], chunk, [&](int g) {
            int k = 0;
            while (g >= off[k + 1]) k++;
            y[k][g - off[k]] = cq_row(*q[k], qx, g - off[k], t);
        });
    }

    Decoder(Model& mm, const std::vector<float>& M, const std::vector<uint8_t>& mv, SpinPool& p) : m(mm), mvalid(mv), pool(p) {
        nM = (int)(M.size() / 512);
        emb = m.Q("embedding/embedding");
        for (int l = 0; l < 8; l++) {
            Layer& y = Ls[l];
            y.q = m.Q(L("self_attn/q_proj/kernel", l));
            y.k = m.Q(L("self_attn/k_proj/kernel", l));
            y.v = m.Q(L("self_attn/v_proj/kernel", l));
            y.gate = m.Q(L("self_attn/gate_proj/kernel", l));
            y.out = m.Q(L("self_attn/out_proj/kernel", l));
            y.cq_ = m.Q(L("cross_attn/q_proj/kernel", l));
            y.cgate = m.Q(L("cross_attn/gate_proj/kernel", l));
            y.cout = m.Q(L("cross_attn/out_proj/kernel", l));
            y.taps[0] = m.F(L("self_attn/q_taps", l));
            y.taps[1] = m.F(L("self_attn/k_taps", l));
            y.taps[2] = m.F(L("self_attn/v_taps", l));
            y.qn = m.F(L("self_attn/q_norm/scale", l));
            y.kn = m.F(L("self_attn/k_norm/scale", l));
            y.pre = m.F(L("ZCRMSNorm_0/scale", l));
            y.post = m.F(L("post_attn_norm/scale", l));
            y.cross = m.F(L("cross_norm/scale", l));
            y.pcross = m.F(L("post_cross_norm/scale", l));
            y.cqn = m.F(L("cross_attn/q_norm/scale", l));
            y.hada = m.F(L("pre_hada_norm/scale", l));
            y.ag = sigmoidf(m.F(L("attn_gate", l))[0]);
            y.cg = sigmoidf(m.F(L("cross_gate", l))[0]);
            auto P = [&](const char* n) { return m.F(L((std::string("hadamard_mlp/") + n).c_str(), l)); };
            y.d1 = P("d1"); y.d2 = P("d2"); y.d3 = P("d3"); y.d4 = P("d4"); y.b2 = P("b2");
            y.cond_v = P("cond_v"); y.cond_u = P("cond_u");
            y.w1a = P("w1a"); y.w1b = P("w1b"); y.w2a = P("w2a"); y.w2b = P("w2b"); y.w3a = P("w3a"); y.w3b = P("w3b");
            ck[l].resize((size_t)nM * 384);
            cv[l].resize((size_t)nM * 512);
            qmm(m, L("cross_attn/k_proj/kernel", l), M.data(), nM, ck[l].data());
            qmm(m, L("cross_attn/v_proj/kernel", l), M.data(), nM, cv[l].data());
            const float* kn = m.F(L("cross_attn/k_norm/scale", l));
            for (int j = 0; j < nM; j++)
                for (int h = 0; h < 8; h++) zc_norm(&ck[l][(size_t)j * 384 + h * 48], 48, kn, &ck[l][(size_t)j * 384 + h * 48]);
#if defined(OASIS_NEON)
            nMp = (nM + 15) / 16 * 16;
            cross->kt[l].assign((size_t)8 * 48 * nMp, 0);
            for (int j = 0; j < nM; j++)
                for (int h = 0; h < 8; h++)
                    for (int c = 0; c < 48; c++)
                        cross->kt[l][(size_t)h * 48 * nMp + (size_t)(j / 16) * 768 + c * 16 + j % 16] = to_half(ck[l][(size_t)j * 384 + h * 48 + c]);
            cross->v16[l].resize((size_t)nM * 512);
            for (int h = 0; h < 8; h++)
                for (int j = 0; j < nM; j++)
                    for (int c = 0; c < 64; c++)
                        cross->v16[l][((size_t)h * nM + j) * 64 + c] = to_half(cv[l][(size_t)j * 512 + h * 64 + c]);
            std::vector<float>().swap(ck[l]);
            std::vector<float>().swap(cv[l]);
#endif
        }
        for (int s = 0; s < 2; s++) {
            std::string pre = "engrams_" + std::to_string(s) + "/";
            sites[s] = {m.Q(pre + "embedding"), m.Q(pre + "key_proj/kernel"), m.Q(pre + "value_proj/kernel"), m.F(pre + "taps")};
        }
        phi_pre = m.Wdq("stack/mhc_phi_pre");
        phi_post = m.Wdq("stack/mhc_phi_post");
        phi_res = m.Wdq("stack/mhc_phi_res");
        a_pre = m.F("stack/mhc_a_pre"); a_post = m.F("stack/mhc_a_post"); a_res = m.F("stack/mhc_a_res");
        b_pre = m.F("stack/mhc_b_pre"); b_post = m.F("stack/mhc_b_post"); b_res = m.F("stack/mhc_b_res");
        final_norm = m.F("stack/final_norm/scale");
    }

    static void kron(const float* in, const float* a, const float* b, float* o) {
#if defined(OASIS_NEON)
        // tmp = in [16, 32] times b [32, 32], then o = a^T [16, 16] times tmp, a row of 32 at a time in
        // registers; every output adds its products in the same order as the loops below.
        float tmp[16 * 32];
        for (int i = 0; i < 16; i++) {
            float32x4_t c[8];
            for (int q = 0; q < 8; q++) c[q] = vdupq_n_f32(0.0f);
            for (int j = 0; j < 32; j++) {
                const float z = in[i * 32 + j];
                const float* br = b + j * 32;
                for (int q = 0; q < 8; q++) c[q] = vfmaq_n_f32(c[q], vld1q_f32(br + 4 * q), z);
            }
            for (int q = 0; q < 8; q++) vst1q_f32(tmp + i * 32 + 4 * q, c[q]);
        }
        for (int k = 0; k < 16; k++) {
            float32x4_t c[8];
            for (int q = 0; q < 8; q++) c[q] = vdupq_n_f32(0.0f);
            for (int i = 0; i < 16; i++) {
                const float av = a[i * 16 + k];
                const float* tr = tmp + i * 32;
                for (int q = 0; q < 8; q++) c[q] = vfmaq_n_f32(c[q], vld1q_f32(tr + 4 * q), av);
            }
            for (int q = 0; q < 8; q++) vst1q_f32(o + k * 32 + 4 * q, c[q]);
        }
#else
        float tmp[16 * 32] = {0};
        for (int i = 0; i < 16; i++)
            for (int j = 0; j < 32; j++) {
                float z = in[i * 32 + j];
                const float* br = b + j * 32;
                for (int l2 = 0; l2 < 32; l2++) tmp[i * 32 + l2] += z * br[l2];
            }
        for (int k = 0; k < 16 * 32; k++) o[k] = 0.0f;
        for (int i = 0; i < 16; i++)
            for (int k = 0; k < 16; k++) {
                float av = a[i * 16 + k];
                for (int l2 = 0; l2 < 32; l2++) o[k * 32 + l2] += av * tmp[i * 32 + l2];
            }
#endif
    }

    void hmlp(const Layer& y, const float* x, float* out) {
        float c8[8], z[512], t[512];
        float mx = -1e30f;
        for (int r = 0; r < 8; r++) c8[r] = 0.0f;
        for (int i = 0; i < 512; i++) {
            const float* row = y.cond_v + (size_t)i * 8;
            for (int r = 0; r < 8; r++) c8[r] += x[i] * row[r];
        }
        for (int r = 0; r < 8; r++) mx = std::max(mx, c8[r]);
        float den = 0;
        for (int r = 0; r < 8; r++) { c8[r] = std::exp(c8[r] - mx); den += c8[r]; }
        float cond[512];
        for (int i = 0; i < 512; i++) cond[i] = 1.0f;
        for (int r = 0; r < 8; r++) {
            float w = c8[r] / den;
            const float* u = y.cond_u + (size_t)r * 512;
            for (int i = 0; i < 512; i++) cond[i] += w * u[i];
        }
        for (int i = 0; i < 512; i++) t[i] = y.d1[i] * x[i];
        kron(t, y.w1a, y.w1b, z);
        for (int i = 0; i < 512; i++) t[i] = z[m.perm1[i]];
#if defined(OASIS_NEON)
        for (int i = 0; i < 512; i++) z[i] = y.d2[i] * cond[i] * t[i] + y.b2[i];
        silu_inplace(z, 512);
#else
        for (int i = 0; i < 512; i++) z[i] = siluf(y.d2[i] * cond[i] * t[i] + y.b2[i]);
#endif
        kron(z, y.w2a, y.w2b, t);
        for (int i = 0; i < 512; i++) z[i] = t[m.perm2[i]];
        for (int i = 0; i < 512; i++) t[i] = y.d3[i] * z[i];
        kron(t, y.w3a, y.w3b, z);
        for (int i = 0; i < 512; i++) out[i] = y.d4[i] * z[i];
    }

    void softmax_rows(float* s, int n) {
        float mx = -1e30f;
        for (int j = 0; j < n; j++) mx = std::max(mx, s[j]);
        float den = 0.0f;
        for (int j = 0; j < n; j++) { s[j] = fexp(s[j] - mx); den += s[j]; }
        float inv = 1.0f / den;
        for (int j = 0; j < n; j++) s[j] *= inv;
    }

    void block(int l, const float* u, int t, float* xout) {
        const Layer& y = Ls[l];
        float x[512], xn[512], tmp[512];
        std::memcpy(x, u, sizeof x);
        int site = l == 3 ? 0 : l == 7 ? 1 : -1;
        if (site >= 0) {
            const Site& st = sites[site];
            float e[512];
            for (int oi = 0; oi < 2; oi++)
                for (int h = 0; h < 2; h++) {
                    int ti = oi * 2 + h;
                    if (t - (ORDERS[oi] - 1) < 0) std::memset(e + ti * 128, 0, 128 * sizeof(float));
                    else cq_dequant_row(st.tab, ti * 18432 + (int)engram_index(toks, t, oi, ORDERS[oi], h), e + ti * 128);
                }
            float k[512], vr[512], v[512] = {0};
            prep_x(e, 512, 2, qx);
            rows({&st.key, &st.val}, {k, vr});
            ev[site].insert(ev[site].end(), vr, vr + 512);
            for (int j = 0; j < 4; j++) {
                int tt = t - 3 * j;
                if (tt < 0) continue;
                const float* tp = st.taps + j * 512;
                const float* src = ev[site].data() + (size_t)tt * 512;
                for (int c = 0; c < 512; c++) v[c] += tp[c] * src[c];
            }
            float ru[512], rk[512];
            rms_unit(x, 512, ru);
            rms_unit(k, 512, rk);
            float alpha = sigmoidf(dotf(ru, rk, 512) / std::sqrt(512.0f));
            for (int c = 0; c < 512; c++) x[c] += alpha * v[c];
        }
        float skip[512];
        std::memcpy(skip, x, sizeof skip);
        zc_norm(x, 512, y.pre, xn);
        prep_x(xn, 512, 2, qx);
        const int dims[3] = {384, 96, 128};
        float qkv[3][384], proj[3][384];
        float gate[512];
        rows({&y.gate, &y.q, &y.k, &y.v}, {gate, proj[0], proj[1], proj[2]});
        for (int s = 0; s < 3; s++) {
            const float* r = proj[s];
            raw[l][s].insert(raw[l][s].end(), r, r + dims[s]);
            const float* taps = y.taps[s];
            for (int c = 0; c < dims[s]; c++) {
                float a = 0;
                for (int j = 0; j < 3; j++)
                    if (t - j >= 0) a += taps[j * dims[s] + c] * raw[l][s][(size_t)(t - j) * dims[s] + c];
                qkv[s][c] = a;
            }
        }
        auto rope48 = [&](float* v) {
            float a[24], b[24];
            std::memcpy(a, v, sizeof a);
            std::memcpy(b, v + 24, sizeof b);
            for (int i = 0; i < 24; i++) {
                float c = rope_c[(size_t)t * 24 + i], s = rope_s[(size_t)t * 24 + i];
                v[i] = a[i] * c - b[i] * s;
                v[24 + i] = b[i] * c + a[i] * s;
            }
        };
        for (int h = 0; h < 8; h++) { zc_norm(qkv[0] + h * 48, 48, y.qn, qkv[0] + h * 48); rope48(qkv[0] + h * 48); }
        for (int h = 0; h < 2; h++) { zc_norm(qkv[1] + h * 48, 48, y.kn, qkv[1] + h * 48); rope48(qkv[1] + h * 48); }
#if defined(OASIS_NEON)
        for (int g = 0; g < 2; g++) {
            K[l][g].insert(K[l][g].end(), qkv[1] + g * 48, qkv[1] + g * 48 + 48);
            V[l][g].insert(V[l][g].end(), qkv[2] + g * 64, qkv[2] + g * 64 + 64);
        }
        const float* kb[2] = {K[l][0].data(), K[l][1].data()};
        const float* vb[2] = {V[l][0].data(), V[l][1].data()};
        const size_t ks = 48, vs = 64;
#else
        K[l].insert(K[l].end(), qkv[1], qkv[1] + 96);
        V[l].insert(V[l].end(), qkv[2], qkv[2] + 128);
#endif
        int n = t + 1;
        float o[512];
        const int span = std::max(n, std::max(nM, nMp));
        scratch.resize((size_t)8 * span);
        const float isq = 1.0f / std::sqrt(48.0f);
        // Heads are independent; each keeps its own score row.
        pool.run(8, n >= 64 ? 1 : 8, [&](int h) {
            float* s = scratch.data() + (size_t)h * span;
            int g = h / 4;
#if defined(OASIS_NEON)
            for (int j = 0; j < n; j++) s[j] = dotf(qkv[0] + h * 48, kb[g] + (size_t)j * ks, 48) * isq;
#else
            for (int j = 0; j < n; j++) s[j] = dotf(qkv[0] + h * 48, &K[l][(size_t)j * 96 + g * 48], 48) * isq;
#endif
            softmax_rows(s, n);
            float* oh = o + h * 64;
            for (int c = 0; c < 64; c++) oh[c] = 0;
            for (int j = 0; j < n; j++) {
#if defined(OASIS_NEON)
                const float* vv = vb[g] + (size_t)j * vs;
#else
                const float* vv = &V[l][(size_t)j * 128 + g * 64];
#endif
                float w = s[j];
                for (int c = 0; c < 64; c++) oh[c] += w * vv[c];
            }
        });
#if defined(OASIS_NEON)
        mul_sigmoid(o, gate, 512);
#else
        for (int c = 0; c < 512; c++) o[c] *= sigmoidf(gate[c]);
#endif
        float A[512];
        prep_x(o, 512, 2, qx);
        rows({&y.out}, {A});
        zc_norm(A, 512, y.post, tmp);
        for (int c = 0; c < 512; c++) x[c] = skip[c] + y.ag * tmp[c];
        std::memcpy(skip, x, sizeof skip);
        float xc[512], qc[384];
        zc_norm(x, 512, y.cross, xc);
        prep_x(xc, 512, 2, qx);
        rows({&y.cq_, &y.cgate}, {qc, gate});
        for (int h = 0; h < 8; h++) zc_norm(qc + h * 48, 48, y.cqn, qc + h * 48);
#if defined(OASIS_NEON)
        // Scores for sixteen memory entries at a time from the transposed fp16 keys, then the fp16 values.
        pool.run(8, nM >= 64 ? 1 : 8, [&](int h) {
            float* s = scratch.data() + (size_t)h * span;
            const uint16_t* kh = cross->kt[l].data() + (size_t)h * 48 * nMp;
            const float* qh = qc + h * 48;
            for (int j0 = 0; j0 < nMp; j0 += 16) {
                float32x4_t a0 = vdupq_n_f32(0.0f), a1 = a0, a2 = a0, a3 = a0;
                for (int c = 0; c < 48; c++) {
                    const uint16_t* kr = kh + (size_t)(j0 / 16) * 768 + c * 16;
                    const float16x8_t k0 = vreinterpretq_f16_u16(vld1q_u16(kr)), k1 = vreinterpretq_f16_u16(vld1q_u16(kr + 8));
                    const float qv = qh[c];
                    a0 = vfmaq_n_f32(a0, vcvt_f32_f16(vget_low_f16(k0)), qv);
                    a1 = vfmaq_n_f32(a1, vcvt_high_f32_f16(k0), qv);
                    a2 = vfmaq_n_f32(a2, vcvt_f32_f16(vget_low_f16(k1)), qv);
                    a3 = vfmaq_n_f32(a3, vcvt_high_f32_f16(k1), qv);
                }
                vst1q_f32(s + j0, a0);
                vst1q_f32(s + j0 + 4, a1);
                vst1q_f32(s + j0 + 8, a2);
                vst1q_f32(s + j0 + 12, a3);
            }
            for (int j = 0; j < nM; j++) s[j] = mvalid[j] ? s[j] * isq : -1e30f;
            softmax_rows(s, nM);
            float32x4_t acc[16];
            for (int q = 0; q < 16; q++) acc[q] = vdupq_n_f32(0.0f);
            const uint16_t* vh = cross->v16[l].data() + (size_t)h * nM * 64;
            for (int j = 0; j < nM; j++) {
                if (!mvalid[j]) continue;
                const float w = s[j];
                const uint16_t* vv = vh + (size_t)j * 64;
                for (int q = 0; q < 8; q++) {
                    const float16x8_t v8 = vreinterpretq_f16_u16(vld1q_u16(vv + q * 8));
                    acc[2 * q] = vfmaq_n_f32(acc[2 * q], vcvt_f32_f16(vget_low_f16(v8)), w);
                    acc[2 * q + 1] = vfmaq_n_f32(acc[2 * q + 1], vcvt_high_f32_f16(v8), w);
                }
            }
            for (int q = 0; q < 16; q++) vst1q_f32(o + h * 64 + q * 4, acc[q]);
        });
#else
        pool.run(8, nM >= 64 ? 1 : 8, [&](int h) {
            float* s = scratch.data() + (size_t)h * span;
            for (int j = 0; j < nM; j++) {
                if (!mvalid[j]) { s[j] = -1e30f; continue; }
                s[j] = dotf(qc + h * 48, &ck[l][(size_t)j * 384 + h * 48], 48) * isq;
            }
            softmax_rows(s, nM);
            float* oh = o + h * 64;
            for (int c = 0; c < 64; c++) oh[c] = 0;
            for (int j = 0; j < nM; j++) {
                if (!mvalid[j]) continue;
                const float* vv = &cv[l][(size_t)j * 512 + h * 64];
                float w = s[j];
                for (int c = 0; c < 64; c++) oh[c] += w * vv[c];
            }
        });
#endif
#if defined(OASIS_NEON)
        mul_sigmoid(o, gate, 512);
#else
        for (int c = 0; c < 512; c++) o[c] *= sigmoidf(gate[c]);
#endif
        prep_x(o, 512, 2, qx);
        rows({&y.cout}, {A});
        zc_norm(A, 512, y.pcross, tmp);
        for (int c = 0; c < 512; c++) x[c] = skip[c] + y.cg * tmp[c];
        zc_norm(x, 512, y.hada, xn);
        hmlp(y, xn, tmp);
        for (int c = 0; c < 512; c++) xout[c] = x[c] + tmp[c];
    }

    std::vector<float> scratch, rope_c, rope_s;

    void ensure_rope(int t) {
        int have = (int)(rope_c.size() / 24);
        if (t < have) return;
        int want = std::max(t + 1, have * 2 + 64);
        rope_c.resize((size_t)want * 24);
        rope_s.resize((size_t)want * 24);
        for (int p = have; p < want; p++)
            for (int i = 0; i < 24; i++) {
                double fr = 1.0 / std::pow(100000.0, (2.0 * i) / 48.0);
                rope_c[(size_t)p * 24 + i] = (float)std::cos(p * fr);
                rope_s[(size_t)p * 24 + i] = (float)std::sin(p * fr);
            }
    }

    // Feed one token; returns logits over the vocabulary.
    std::vector<float> step(int tok) {
        int t = (int)toks.size();
        toks.push_back(tok);
        ensure_rope(t);
        float X[4][512], row[512];
        cq_dequant_row(emb, tok, row);
        for (int ln = 0; ln < 4; ln++)
            for (int c = 0; c < 512; c++) X[ln][c] = std::sqrt(512.0f) * row[c];
        for (int l = 0; l < 8; l++) {
            float nx[2048];
            rms_unit(&X[0][0], 2048, nx);
            float hv[24];
#if defined(OASIS_NEON)
            const int phi_chunk = 8;   // a contiguous share of the projections for each of three threads
            pool.run(24, phi_chunk, [&](int i) {
#else
            pool.run(24, 4, [&](int i) {
#endif
                const float* row = i < 4 ? phi_pre + (size_t)(4 * l + i) * 2048
                                 : i < 8 ? phi_post + (size_t)(4 * l + i - 4) * 2048
                                         : phi_res + (size_t)(16 * l + i - 8) * 2048;
                hv[i] = dotf(row, nx, 2048);
            });
            float hp[4], hq[4], hr[16];
            for (int i = 0; i < 4; i++) {
                hp[i] = hv[i];
                hq[i] = hv[4 + i];
            }
            for (int i = 0; i < 16; i++) hr[i] = hv[8 + i];
            float u[512] = {0}, yv[512];
            for (int i = 0; i < 4; i++) {
                float off = (l % 4 == i) ? 4.0f : -4.0f;
                hp[i] = sigmoidf(a_pre[l] * hp[i] + b_pre[l * 4 + i] + off);
                hq[i] = 2.0f * sigmoidf(a_post[l] * hq[i] + b_post[l * 4 + i] + ((l % 4 == i) ? 0.0f : -4.0f));
                for (int c = 0; c < 512; c++) u[c] += hp[i] * X[i][c];
            }
            for (int i = 0; i < 16; i++) hr[i] = a_res[l] * hr[i] + b_res[l * 16 + i];
            sinkhorn(hr);
            block(l, u, t, yv);
            for (int c = 0; c < 512; c++) yv[c] -= u[c];
            float Xn[4][512];
            for (int i = 0; i < 4; i++)
                for (int c = 0; c < 512; c++)
                    Xn[i][c] = hq[i] * yv[c] + hr[i * 4] * X[0][c] + hr[i * 4 + 1] * X[1][c] + hr[i * 4 + 2] * X[2][c] + hr[i * 4 + 3] * X[3][c];
            std::memcpy(X, Xn, sizeof X);
        }
        float h[512], hn[512];
        for (int c = 0; c < 512; c++) h[c] = 0.25f * (X[0][c] + X[1][c] + X[2][c] + X[3][c]);
        zc_norm(h, 512, final_norm, hn);
        std::vector<float> logits(emb.rows);
        prep_x(hn, 512, emb.bits, qx);
        rows({&emb}, {logits.data()}, 128);
        return logits;
    }
};

static int decode_threads() { return std::max(1, std::min(engine_threads(), 8)); }

// Keeps the pool's workers spinning for the length of one decode.
struct Awake {
    SpinPool& p;
    explicit Awake(SpinPool& pool) : p(pool) { p.wake(); }
    ~Awake() { p.sleep(); }
};

// When the sequence ends in one short block written several times over, which the decoder only does
// once it has fallen into a loop (a block of 2 to 8 tokens four times, or one token six times), returns
// how many trailing tokens repeat the first copy; otherwise 0.
static int looping(const std::vector<int>& s) {
    const int n = (int)s.size();
    for (int p = 1; p <= 8; p++) {
        int reps = p == 1 ? 6 : 4;
        if (n < p * reps) continue;
        bool same = true;
        for (int k = 1; k < reps && same; k++)
            for (int i = 0; i < p && same; i++) same = s[n - 1 - i] == s[n - 1 - i - k * p];
        if (same) return (reps - 1) * p;
    }
    return 0;
}

// Longest reasoning trace kept; a trace that has not closed by then is closed by the engine.
static const int TRACE_MAX = 256;

// Keyword biasing: a small logit bonus for a keyword's opening word piece, and a larger one for each
// following piece once the earlier ones have been written, so a keyword the model starts is completed
// without common word starts being pushed into every line.
static const float KEYWORD_START_BONUS = 1.0f;
static const float KEYWORD_NEXT_BONUS = 3.0f;

std::vector<int> generate(Model& m, const std::vector<float>& M, const std::vector<uint8_t>& mvalid,
                          const std::vector<int>& refs, const std::string& language, int max_tokens,
                          const std::vector<std::string>& keywords) {
    std::lock_guard<std::mutex> lock(m.decode_mutex);
    if (!m.pool) m.pool = std::make_shared<SpinPool>(decode_threads());
    Decoder dec(m, M, mvalid, *m.pool);
    Awake awake(*m.pool);
    std::vector<uint8_t> banned(m.vocab, 0);
    for (auto& kv : m.langs) banned[kv.second] = 1;
    for (int i = m.ref_base; i < m.vocab; i++) banned[i] = 1;
    for (int r : refs) banned[r] = 0;
    auto it = m.langs.find(language);
    int lang = it != m.langs.end() ? it->second : m.langs.at("en");
    // A keyword follows a space in names and the trace, or "=" in attribute values; each spelling is its
    // own token sequence. Only the spaced spelling, which opens on a word start, gets the opening bonus.
    std::vector<std::pair<std::vector<int>, bool>> seqs;
    for (const auto& k : keywords)
        for (int spaced = 1; spaced >= 0; spaced--) {
            std::vector<int> s = encode_text(m, spaced ? " " + k : k);
            bool dup = false;
            for (const auto& q : seqs) dup = dup || q.first == s;
            if (!s.empty() && !dup) seqs.push_back({s, spaced == 1});
        }
    int think = -1, think_end = -1;
    for (const auto& mk : m.markers) {
        if (mk.first == "<think>") think = mk.second;
        if (mk.first == "</think>") think_end = mk.second;
    }
    std::vector<float> logits = dec.step(m.bos);
    logits = dec.step(lang);
    std::vector<int> out;
    std::vector<std::pair<int, float>> bonus;
    bool in_trace = false;
    int trace_len = 0;
    for (int s = 0; s < max_tokens; s++) {
        bonus.clear();
        for (const auto& [q, opens] : seqs) {
            int n = std::min((int)q.size() - 1, (int)out.size());
            for (int len = n; len >= 0; len--)
                if (std::equal(q.begin(), q.begin() + len, out.end() - len)) {
                    if (!len && !opens) break;
                    float add = len ? KEYWORD_NEXT_BONUS : KEYWORD_START_BONUS;
                    bool have = false;
                    for (auto& b : bonus)
                        if (b.first == q[len]) { b.second = std::max(b.second, add); have = true; }
                    if (!have) bonus.push_back({q[len], add});
                    break;
                }
        }
        for (const auto& b : bonus) logits[b.first] += b.second;
        int best = -1;
        float bv = -1e30f;
        for (int i = 0; i < (int)logits.size(); i++)
            if (!banned[i] && logits[i] > bv) { bv = logits[i]; best = i; }
        if (best < 0 || best == m.eos) break;
        // A trace that loops or runs on is closed so that the description still gets written; a
        // description that loops ends where the loop began to repeat. Repeated copies are dropped.
        out.push_back(best);
        if (best == think && think >= 0) {
            in_trace = true;
        } else if (in_trace) {
            int rep = looping(out);
            if (best == think_end) {
                in_trace = false;
            } else if (++trace_len > TRACE_MAX || rep) {
                out.resize(out.size() - std::max(rep, 1));   // the newest token was never fed to the decoder
                out.push_back(think_end);
                in_trace = false;
            }
        } else if (int rep = looping(out)) {
            out.resize(out.size() - rep);
            break;
        }
        logits = dec.step(out.back());
    }
    return out;
}

// BPE over one marker-free segment: adjacent symbols merged lowest piece id first (piece scores fall
// with id, so this is SentencePiece's highest-score-first order), unknown symbols as bytes.
static void bpe(const Model& m, const std::string& seg, std::vector<int>& ids) {
    std::vector<std::string> syms;
    for (size_t i = 0; i < seg.size();) {
        unsigned char c = (unsigned char)seg[i];
        size_t n = c < 0x80 ? 1 : (c >> 5) == 6 ? 2 : (c >> 4) == 14 ? 3 : (c >> 3) == 30 ? 4 : 1;
        if (i + n > seg.size()) n = 1;
        syms.push_back(seg.substr(i, n));
        i += n;
    }
    for (;;) {
        int best = INT_MAX;
        size_t at = SIZE_MAX;
        for (size_t i = 0; i + 1 < syms.size(); i++) {
            auto it = m.piece_id.find(syms[i] + syms[i + 1]);
            if (it != m.piece_id.end() && it->second < best) { best = it->second; at = i; }
        }
        if (at == SIZE_MAX) break;
        syms[at] += syms[at + 1];
        syms.erase(syms.begin() + (long)at + 1);
    }
    for (const auto& s : syms) {
        auto it = m.piece_id.find(s);
        if (it != m.piece_id.end()) ids.push_back(it->second);
        else for (unsigned char b : s) ids.push_back(m.byte_id[b]);
    }
}

// SentencePiece encoding as the vocabulary specifies it: spaces become U+2581, no dummy prefix,
// user-defined pieces such as <think> or language tags taken whole, the text between them by BPE.
std::vector<int> encode_text(const Model& m, const std::string& text) {
    std::string esc;
    for (char c : text) {
        if (c == ' ') esc += "\xE2\x96\x81";
        else esc += c;
    }
    std::vector<int> ids;
    size_t i = 0, start = 0;
    while (i < esc.size()) {
        const std::pair<std::string, int>* hit = nullptr;
        for (const auto& mk : m.markers)
            if (esc.compare(i, mk.first.size(), mk.first) == 0) { hit = &mk; break; }
        if (!hit) { i++; continue; }
        bpe(m, esc.substr(start, i - start), ids);
        ids.push_back(hit->second);
        i += hit->first.size();
        start = i;
    }
    bpe(m, esc.substr(start), ids);
    return ids;
}

static std::string json_str(const std::string& s) {
    std::string o = "\"";
    for (unsigned char c : s) {
        if (c == '"') o += "\\\"";
        else if (c == '\\') o += "\\\\";
        else if (c == '\n') o += "\\n";
        else if (c < 0x20) { char b[8]; std::snprintf(b, sizeof b, "\\u%04x", c); o += b; }
        else o += (char)c;
    }
    return o + "\"";
}

static std::string json_num(double v) {
    if (!std::isfinite(v)) return "null";
    char b[32];
    std::snprintf(b, sizeof b, "%.6g", v);
    return b;
}

// Log-probability of tok under the softmax over the tokens not in banned.
static double log_prob(const std::vector<float>& L, const std::vector<uint8_t>& banned, int tok) {
    if (banned[tok]) return -INFINITY;
    float mx = -1e30f;
    for (size_t i = 0; i < L.size(); i++) if (!banned[i] && L[i] > mx) mx = L[i];
    double s = 0.0;
    for (size_t i = 0; i < L.size(); i++) if (!banned[i]) s += std::exp((double)L[i] - mx);
    return (double)L[tok] - mx - std::log(s);
}

// A form is one request line naming every field, then one answer line per field:
//   ? sign=stop|yield; n#0..3; note$60; where@sign
//   sign=stop
//   n=1 ...
// Choices and counts are scored as complete answers with their terminator (a newline, or </s> after the
// last field), choices with "-" for abstention; text is decoded greedily; a box is a softmax over the
// occupied reference tokens. Each answer is written into the context before the next field.
std::string fill_form(Model& m, const std::vector<float>& M, const std::vector<uint8_t>& mvalid, const std::vector<int>& refs,
                      const std::vector<std::vector<int>>& ref_boxes, const std::string& language, const std::vector<FormField>& form) {
    std::lock_guard<std::mutex> lock(m.decode_mutex);
    if (!m.pool) m.pool = std::make_shared<SpinPool>(decode_threads());
    Decoder dec(m, M, mvalid, *m.pool);
    Awake awake(*m.pool);
    std::vector<uint8_t> banned(m.vocab, 0), no_refs(m.vocab, 0);
    for (auto& kv : m.langs) banned[kv.second] = no_refs[kv.second] = 1;
    for (int i = m.ref_base; i < m.vocab; i++) banned[i] = no_refs[i] = 1;
    for (int r : refs) banned[r] = 0;
    auto lit = m.langs.find(language);
    int lang = lit != m.langs.end() ? lit->second : m.langs.at("en");

    std::string line = "? ";
    for (size_t i = 0; i < form.size(); i++) {
        const FormField& f = form[i];
        if (i) line += "; ";
        line += f.name;
        if (f.kind == FormField::CHOICE) {
            line += "=";
            for (size_t c = 0; c < f.choices.size(); c++) line += (c ? "|" : "") + f.choices[c];
        } else if (f.kind == FormField::COUNT) {
            line += "#" + std::to_string(f.lo) + ".." + std::to_string(f.hi);
        } else if (f.kind == FormField::TEXT) {
            line += "$" + std::to_string(f.max_length);
        } else {
            line += "@" + f.of;
        }
    }
    std::vector<int> prompt = {m.bos, lang};
    for (int t : encode_text(m, line + "\n")) prompt.push_back(t);
    if (prompt.size() > 1024) throw std::runtime_error("form exceeds the decoder context");
    std::vector<float> logits;
    for (int t : prompt) logits = dec.step(t);
    const std::vector<int> newline = encode_text(m, "\n");

    std::ostringstream js;
    js << "{";
    for (size_t i = 0; i < form.size(); i++) {
        const FormField& f = form[i];
        const bool last = i + 1 == form.size();
        const std::vector<int> term = last ? std::vector<int>{m.eos} : newline;
        for (int t : encode_text(m, f.name + "=")) logits = dec.step(t);
        js << (i ? "," : "") << json_str(f.name) << ":";
        std::vector<int> answer;   // tokens written into the context after "name="
        if (f.kind == FormField::CHOICE || f.kind == FormField::COUNT) {
            std::vector<std::string> cands;
            if (f.kind == FormField::CHOICE) {
                cands = f.choices;
                cands.push_back("-");
            } else {
                for (long v = f.lo; v <= f.hi; v++) cands.push_back(std::to_string(v));
            }
            std::vector<std::vector<int>> toks(cands.size());
            std::vector<double> score(cands.size());
            for (size_t c = 0; c < cands.size(); c++) {
                toks[c] = encode_text(m, cands[c]);
                toks[c].insert(toks[c].end(), term.begin(), term.end());
                Decoder d = dec;
                std::vector<float> L = logits;
                double s = 0.0;
                for (size_t k = 0; k < toks[c].size(); k++) {
                    s += log_prob(L, banned, toks[c][k]);
                    if (k + 1 < toks[c].size()) L = d.step(toks[c][k]);
                }
                score[c] = s;
            }
            double mx = -INFINITY, den = 0.0;
            size_t best = 0;
            for (size_t c = 0; c < cands.size(); c++) if (score[c] > mx) { mx = score[c]; best = c; }
            for (double s : score) den += std::exp(s - mx);
            if (f.kind == FormField::CHOICE)
                js << "{\"value\":" << (best == f.choices.size() ? "null" : f.values[best]);
            else
                js << "{\"value\":" << (f.lo + (long)best);
            js << ",\"probabilities\":{";
            for (size_t c = 0; c < cands.size(); c++)
                js << (c ? "," : "") << json_str(cands[c]) << ":" << json_num(std::exp(score[c] - mx) / den);
            js << "}}";
            answer = toks[best];
        } else if (f.kind == FormField::TEXT) {
            std::vector<int> gen;
            double s = 0.0;
            bool ended = false;
            for (int step = 0; step < 4 * f.max_length + 8; step++) {
                int tok = -1;
                float bv = -1e30f;
                for (int v = 0; v < (int)logits.size(); v++)
                    if (!no_refs[v] && logits[v] > bv) { bv = logits[v]; tok = v; }
                s += log_prob(logits, no_refs, tok);
                if (tok == term[0] || tok == m.eos || tok == newline[0]) { ended = true; break; }
                gen.push_back(tok);
                logits = dec.step(tok);
                std::string sofar = detokenize(m, gen);
                size_t chars = 0;
                for (unsigned char c : sofar) chars += (c & 0xC0) != 0x80;
                if ((int)chars >= f.max_length) break;
            }
            if (!ended) s += log_prob(logits, no_refs, term[0]);
            js << "{\"value\":" << json_str(detokenize(m, gen)) << ",\"probability\":" << json_num(std::exp(s)) << "}";
            answer = term;   // the generated tokens are already in the context
        } else {
            if (refs.empty()) {
                js << "{\"value\":null,\"probability\":null}";
                answer = encode_text(m, "-");
                answer.insert(answer.end(), term.begin(), term.end());
            } else {
                float mx = -1e30f;
                for (int r : refs) mx = std::max(mx, logits[r]);
                double den = 0.0;
                size_t best = 0;
                for (size_t k = 0; k < refs.size(); k++) {
                    den += std::exp((double)logits[refs[k]] - mx);
                    if (logits[refs[k]] > logits[refs[best]]) best = k;
                }
                const std::vector<int>& b = ref_boxes[best];
                js << "{\"value\":[" << b[0] << "," << b[1] << "," << b[2] << "," << b[3] << "],\"probability\":"
                   << json_num(std::exp((double)logits[refs[best]] - mx) / den) << "}";
                answer = {refs[best]};
                answer.insert(answer.end(), term.begin(), term.end());
            }
        }
        if (!last)
            for (int t : answer) logits = dec.step(t);
    }
    js << "}";
    return js.str();
}

std::string detokenize(const Model& m, const std::vector<int>& ids) {
    std::string out, bytes;
    auto flush = [&]() { out += bytes; bytes.clear(); };
    for (int id : ids) {
        if (id >= m.n_pieces) {
            flush();
            int k = id - m.ref_base;
            out += k < 64 ? "<obj" + std::to_string(k) + ">" : "<txt" + std::to_string(k - 64) + ">";
            continue;
        }
        const std::string& p = m.pieces[id];
        int t = m.piece_types[id];
        if (t == 4 && p.size() == 6 && p.rfind("<0x", 0) == 0) {
            bytes += (char)std::stoi(p.substr(3, 2), nullptr, 16);
            continue;
        }
        flush();
        if (t == 0) {
            std::string s = p;
            size_t pos;
            while ((pos = s.find("\xE2\x96\x81")) != std::string::npos) s.replace(pos, 3, " ");
            out += s;
        } else {
            out += p;
        }
    }
    flush();
    return out;
}

}  // namespace oasis

// ==== Record assembly, mask outlines and bitmaps, forms, options parsing and the C API

namespace oasis {

// The calling thread's OpenMP thread count set to engine_threads() for one call, and restored after it.
struct ThreadScope {
#ifdef _OPENMP
    int prev = omp_get_max_threads();
    ThreadScope() { omp_set_num_threads(engine_threads()); }
    ~ThreadScope() { omp_set_num_threads(prev); }
#endif
};

// ------------------------------------------------------------------ JSON writing
static std::string esc(const std::string& s) {
    std::string o = "\"";
    for (unsigned char c : s) {
        if (c == '"') o += "\\\"";
        else if (c == '\\') o += "\\\\";
        else if (c == '\n') o += "\\n";
        else if (c == '\t') o += "\\t";
        else if (c < 0x20) { char b[8]; std::snprintf(b, sizeof b, "\\u%04x", c); o += b; }
        else o += (char)c;
    }
    return o + "\"";
}

static std::string num(double v, int prec = 3) {
    if (!std::isfinite(v)) return "null";
    char b[32];
    std::snprintf(b, sizeof b, "%.*f", prec, v);
    std::string s = b;
    if (s.find('.') != std::string::npos) {
        while (s.back() == '0') s.pop_back();
        if (s.back() == '.') s.pop_back();
    }
    return s;
}

static std::string trim(const std::string& s) {
    size_t a = s.find_first_not_of(" \t\r"), b = s.find_last_not_of(" \t\r");
    return a == std::string::npos ? "" : s.substr(a, b - a + 1);
}

// ------------------------------------------------------------------ JSON reading
struct JV {
    enum Type { NUL, BOOL, NUM, STR, ARR, OBJ } t = NUL;
    bool b = false;
    double n = 0.0;
    std::string s;   // string value, or a number's source text
    std::vector<JV> a;
    std::vector<std::pair<std::string, JV>> o;   // members in document order
    const JV* get(const char* k) const {
        for (const auto& kv : o)
            if (kv.first == k) return &kv.second;
        return nullptr;
    }
};

struct JReader {
    const char* p;
    const char* e;
    [[noreturn]] void fail(const std::string& what) { throw std::runtime_error("invalid options: " + what); }
    void ws() {
        while (p < e && (*p == ' ' || *p == '\t' || *p == '\n' || *p == '\r')) p++;
    }
    static void utf8(std::string& out, unsigned cp) {
        if (cp < 0x80) out += (char)cp;
        else if (cp < 0x800) { out += (char)(0xC0 | (cp >> 6)); out += (char)(0x80 | (cp & 0x3F)); }
        else if (cp < 0x10000) { out += (char)(0xE0 | (cp >> 12)); out += (char)(0x80 | ((cp >> 6) & 0x3F)); out += (char)(0x80 | (cp & 0x3F)); }
        else { out += (char)(0xF0 | (cp >> 18)); out += (char)(0x80 | ((cp >> 12) & 0x3F)); out += (char)(0x80 | ((cp >> 6) & 0x3F)); out += (char)(0x80 | (cp & 0x3F)); }
    }
    unsigned hex4() {
        if (e - p < 4) fail("bad \\u escape");
        unsigned v = 0;
        for (int i = 0; i < 4; i++) {
            char c = *p++;
            v <<= 4;
            if (c >= '0' && c <= '9') v |= (unsigned)(c - '0');
            else if (c >= 'a' && c <= 'f') v |= (unsigned)(c - 'a' + 10);
            else if (c >= 'A' && c <= 'F') v |= (unsigned)(c - 'A' + 10);
            else fail("bad \\u escape");
        }
        return v;
    }
    std::string str() {
        p++;
        std::string out;
        while (p < e && *p != '"') {
            char c = *p++;
            if (c != '\\') { out += c; continue; }
            if (p >= e) fail("unterminated string");
            char x = *p++;
            switch (x) {
                case '"': out += '"'; break;
                case '\\': out += '\\'; break;
                case '/': out += '/'; break;
                case 'b': out += '\b'; break;
                case 'f': out += '\f'; break;
                case 'n': out += '\n'; break;
                case 'r': out += '\r'; break;
                case 't': out += '\t'; break;
                case 'u': {
                    unsigned cp = hex4();
                    if (cp >= 0xD800 && cp < 0xDC00 && e - p >= 6 && p[0] == '\\' && p[1] == 'u') {
                        p += 2;
                        unsigned lo = hex4();
                        cp = 0x10000 + ((cp - 0xD800) << 10) + (lo - 0xDC00);
                    }
                    utf8(out, cp);
                    break;
                }
                default: fail("bad escape in string");
            }
        }
        if (p >= e) fail("unterminated string");
        p++;
        return out;
    }
    JV value() {
        ws();
        if (p >= e) fail("unexpected end of JSON");
        JV v;
        if (*p == '{') {
            p++;
            v.t = JV::OBJ;
            ws();
            if (p < e && *p == '}') { p++; return v; }
            for (;;) {
                ws();
                if (p >= e || *p != '"') fail("expected a key");
                std::string k = str();
                ws();
                if (p >= e || *p != ':') fail("expected ':'");
                p++;
                JV x = value();
                v.o.emplace_back(std::move(k), std::move(x));
                ws();
                if (p < e && *p == ',') { p++; continue; }
                if (p < e && *p == '}') { p++; return v; }
                fail("expected ',' or '}'");
            }
        }
        if (*p == '[') {
            p++;
            v.t = JV::ARR;
            ws();
            if (p < e && *p == ']') { p++; return v; }
            for (;;) {
                v.a.push_back(value());
                ws();
                if (p < e && *p == ',') { p++; continue; }
                if (p < e && *p == ']') { p++; return v; }
                fail("expected ',' or ']'");
            }
        }
        if (*p == '"') {
            v.t = JV::STR;
            v.s = str();
            return v;
        }
        auto word = [&](const char* w) {
            size_t n = std::strlen(w);
            if ((size_t)(e - p) >= n && std::strncmp(p, w, n) == 0) { p += n; return true; }
            return false;
        };
        if (word("true")) { v.t = JV::BOOL; v.b = true; return v; }
        if (word("false")) { v.t = JV::BOOL; return v; }
        if (word("null")) return v;
        const char* q = p;
        while (p < e && (std::isdigit((unsigned char)*p) || *p == '-' || *p == '+' || *p == '.' || *p == 'e' || *p == 'E')) p++;
        if (q == p) fail("unexpected character");
        v.t = JV::NUM;
        v.s.assign(q, p);
        char* end = nullptr;
        v.n = std::strtod(v.s.c_str(), &end);
        if (!end || *end) fail("bad number " + v.s);
        return v;
    }
};

// ------------------------------------------------------------------ mask geometry
using Ring = std::vector<std::pair<float, float>>;   // vertices in mask-grid units, cell (i, j) centred at (j + 0.5, i + 0.5)

static double ring_area(const Ring& r) {
    double a = 0.0;
    for (size_t i = 0, n = r.size(); i < n; i++) {
        const auto& p = r[i];
        const auto& q = r[(i + 1) % n];
        a += (double)p.first * q.second - (double)q.first * p.second;
    }
    return 0.5 * a;
}

// Douglas-Peucker on a closed ring: split at the vertex farthest from the first, simplify both halves.
static Ring simplify_ring(const Ring& pts, float tol) {
    const size_t n = pts.size();
    if (n < 4) return pts;
    size_t far = 0;
    float fd = -1.0f;
    for (size_t i = 1; i < n; i++) {
        float d = std::hypot(pts[i].first - pts[0].first, pts[i].second - pts[0].second);
        if (d > fd) { fd = d; far = i; }
    }
    std::vector<uint8_t> keep(n, 0);
    keep[0] = keep[far] = 1;
    auto pt = [&](size_t i) { return pts[i % n]; };
    std::vector<std::pair<size_t, size_t>> st{{0, far}, {far, n}};
    while (!st.empty()) {
        auto [a, c] = st.back();
        st.pop_back();
        auto A = pt(a), C = pt(c);
        float L = std::hypot(C.first - A.first, C.second - A.second), best = 0.0f;
        size_t bi = 0;
        for (size_t i = a + 1; i < c; i++) {
            auto P = pt(i);
            float d = L > 0 ? std::fabs((C.first - A.first) * (A.second - P.second) - (A.first - P.first) * (C.second - A.second)) / L
                            : std::hypot(P.first - A.first, P.second - A.second);
            if (d > best) { best = d; bi = i; }
        }
        if (bi && best > tol) { keep[bi % n] = 1; st.push_back({a, bi}); st.push_back({bi, c}); }
    }
    Ring out;
    for (size_t i = 0; i < n; i++)
        if (keep[i]) out.push_back(pts[i]);
    return out;
}

// Exterior outlines of the thresholded 128 x 128 probability field, largest first. Marching squares over
// the cell centres, with every crossing interpolated linearly between the two neighbouring probabilities;
// the field is zero outside the grid, so every outline closes. Hole boundaries are left out.
static std::vector<Ring> outlines(const std::vector<float>& prob, float threshold, bool largest_only) {
    const int N = 128, M = N + 2;   // lattice -1 .. N, stored with an offset of one
    const float t = std::max(threshold, 1e-6f);
    auto f = [&](int i, int j) -> float {
        if (i < 0 || j < 0 || i >= N || j >= N) return -t;
        return prob[(size_t)i * N + j] - t;
    };
    auto inside = [&](int i, int j) { return f(i, j) >= 0.0f; };
    // Edge ids: horizontal edge from lattice (i, j) to (i, j + 1), vertical edge from (i, j) to (i + 1, j).
    auto hid = [&](int i, int j) { return (((i + 1) * M) + (j + 1)) * 2; };
    auto vid = [&](int i, int j) { return (((i + 1) * M) + (j + 1)) * 2 + 1; };
    auto cross = [&](int i0, int j0, int i1, int j1) {
        float a = f(i0, j0), b = f(i1, j1);
        float s = a / (a - b);
        return std::make_pair(j0 + 0.5f + s * (j1 - j0), i0 + 0.5f + s * (i1 - i0));
    };
    enum { T, R, B, L };
    struct Seg { int from, to; std::pair<float, float> p; };
    std::vector<Seg> segs;
    std::vector<int> start_at((size_t)M * M * 2, -1);
    for (int i = -1; i < N; i++)
        for (int j = -1; j < N; j++) {
            int code = (inside(i, j) ? 8 : 0) | (inside(i, j + 1) ? 4 : 0) | (inside(i + 1, j + 1) ? 2 : 0) | (inside(i + 1, j) ? 1 : 0);
            if (code == 0 || code == 15) continue;
            const int eid[4] = {hid(i, j), vid(i, j + 1), hid(i + 1, j), vid(i, j)};
            auto point = [&](int e) {
                switch (e) {
                    case T: return cross(i, j, i, j + 1);
                    case R: return cross(i, j + 1, i + 1, j + 1);
                    case B: return cross(i + 1, j, i + 1, j + 1);
                    default: return cross(i, j, i + 1, j);
                }
            };
            // Directed so that the inside lies to the right of travel on the image grid (y down).
            int pairs[2][2], np = 0;
            auto add = [&](int a, int b) { pairs[np][0] = a; pairs[np][1] = b; np++; };
            float centre = 0.25f * (f(i, j) + f(i, j + 1) + f(i + 1, j + 1) + f(i + 1, j));
            switch (code) {
                case 1: add(L, B); break;
                case 2: add(B, R); break;
                case 3: add(L, R); break;
                case 4: add(R, T); break;
                case 6: add(B, T); break;
                case 7: add(L, T); break;
                case 8: add(T, L); break;
                case 9: add(T, B); break;
                case 11: add(T, R); break;
                case 12: add(R, L); break;
                case 13: add(R, B); break;
                case 14: add(B, L); break;
                case 5:
                    if (centre >= 0.0f) { add(L, T); add(R, B); } else { add(R, T); add(L, B); }
                    break;
                case 10:
                    if (centre >= 0.0f) { add(T, R); add(B, L); } else { add(T, L); add(B, R); }
                    break;
            }
            for (int k = 0; k < np; k++) {
                start_at[eid[pairs[k][0]]] = (int)segs.size();
                segs.push_back({eid[pairs[k][0]], eid[pairs[k][1]], point(pairs[k][0])});
            }
        }
    std::vector<uint8_t> used(segs.size(), 0);
    std::vector<Ring> rings;
    for (size_t s0 = 0; s0 < segs.size(); s0++) {
        if (used[s0]) continue;
        Ring r;
        int s = (int)s0;
        while (s >= 0 && !used[s]) {
            used[s] = 1;
            r.push_back(segs[s].p);
            s = start_at[segs[s].to];
        }
        if (r.size() >= 3) rings.push_back(std::move(r));
    }
    if (rings.empty()) return rings;
    // Outer boundaries share one orientation and holes the other; the largest ring is always an outer one.
    size_t big = 0;
    for (size_t k = 1; k < rings.size(); k++)
        if (std::fabs(ring_area(rings[k])) > std::fabs(ring_area(rings[big]))) big = k;
    const bool positive = ring_area(rings[big]) > 0;
    std::vector<std::pair<double, Ring>> outer;
    for (auto& r : rings) {
        double a = ring_area(r);
        if ((a > 0) == positive && std::fabs(a) >= 0.25) outer.push_back({std::fabs(a), simplify_ring(r, 0.5f)});
    }
    std::stable_sort(outer.begin(), outer.end(), [](const auto& a, const auto& b) { return a.first > b.first; });
    std::vector<Ring> out;
    for (auto& o : outer) {
        if (o.second.size() < 3) continue;
        out.push_back(std::move(o.second));
        if (largest_only) break;
    }
    return out;
}

// All pieces joined into one path: each piece is entered from the nearest vertex of the path so far,
// walked once round, and left by the same bridge.
static Ring merged_ring(const std::vector<Ring>& rings) {
    if (rings.empty()) return {};
    Ring path = rings[0];
    for (size_t k = 1; k < rings.size(); k++) {
        const Ring& r = rings[k];
        size_t bi = 0, bj = 0;
        float bd = 1e30f;
        for (size_t i = 0; i < path.size(); i++)
            for (size_t j = 0; j < r.size(); j++) {
                float dx = path[i].first - r[j].first, dy = path[i].second - r[j].second, d = dx * dx + dy * dy;
                if (d < bd) { bd = d; bi = i; bj = j; }
            }
        Ring out(path.begin(), path.begin() + (long)bi + 1);
        for (size_t t = 0; t <= r.size(); t++) out.push_back(r[(bj + t) % r.size()]);
        out.insert(out.end(), path.begin() + (long)bi, path.end());
        path.swap(out);
    }
    return path;
}

// Column-major run-length encoding at the original image size, starting with a background run: the
// probability field is sampled bilinearly at every pixel centre, thresholded and clipped to the box.
static std::string bitmap_rle(const std::vector<float>& prob, const int bb[4], const Meta& mt, float threshold) {
    const int W = mt.W, H = mt.H;
    const double bx0 = bb[0] / 1000.0 * W, by0 = bb[1] / 1000.0 * H, bx1 = bb[2] / 1000.0 * W, by1 = bb[3] / 1000.0 * H;
    const int x0 = std::max(0, (int)std::floor(bx0)), x1 = std::min(W, (int)std::ceil(bx1));
    const int y0 = std::max(0, (int)std::floor(by0)), y1 = std::min(H, (int)std::ceil(by1));
    std::vector<long long> counts;
    long long run = 0;
    bool fg = false;
    auto push = [&](bool v, long long n) {
        if (n <= 0) return;
        if (v == fg) { run += n; return; }
        counts.push_back(run);
        run = n;
        fg = v;
    };
    const double k = mt.r / 4.0;   // original pixels -> mask cells
    for (int x = 0; x < W; x++) {
        if (x < x0 || x >= x1 || y1 <= y0) { push(false, H); continue; }
        push(false, y0);
        double u = (x + 0.5) * k - 0.5;
        int j0 = (int)std::floor(u);
        double fu = u - j0;
        int ja = std::clamp(j0, 0, 127), jb = std::clamp(j0 + 1, 0, 127);
        bool colin = x + 0.5 >= bx0 && x + 0.5 <= bx1;
        for (int y = y0; y < y1; y++) {
            bool v = false;
            if (colin && y + 0.5 >= by0 && y + 0.5 <= by1) {
                double w = (y + 0.5) * k - 0.5;
                int i0 = (int)std::floor(w);
                double fv = w - i0;
                int ia = std::clamp(i0, 0, 127), ib = std::clamp(i0 + 1, 0, 127);
                double p = (1 - fv) * ((1 - fu) * prob[ia * 128 + ja] + fu * prob[ia * 128 + jb]) +
                           fv * ((1 - fu) * prob[ib * 128 + ja] + fu * prob[ib * 128 + jb]);
                v = p >= threshold;
            }
            push(v, 1);
        }
        push(false, H - y1);
    }
    counts.push_back(run);
    std::ostringstream os;
    os << "{\"size\":[" << H << "," << W << "],\"counts\":[";
    for (size_t i = 0; i < counts.size(); i++) os << (i ? "," : "") << counts[i];
    os << "]}";
    return os.str();
}

// ------------------------------------------------------------------ describe
struct Entity {
    int n = 0;
    std::string name, position, count;
    std::vector<std::pair<std::string, std::string>> attrs;
    std::vector<std::pair<char, int>> refs;
};

static std::vector<std::pair<std::string, std::string>> parse_attrs(const std::string& s) {
    std::vector<std::pair<std::string, std::string>> out;
    std::stringstream ss(s);
    std::string part;
    while (std::getline(ss, part, ';')) {
        auto eq = part.find('=');
        if (eq == std::string::npos) continue;
        out.push_back({trim(part.substr(0, eq)), trim(part.substr(eq + 1))});
    }
    return out;
}

static bool has(const std::vector<std::string>& v, const char* x) { return std::find(v.begin(), v.end(), x) != v.end(); }

std::string describe_json(Model& m, const uint8_t* rgb, int W, int H, const Options& o, double* tm) {
    ThreadScope threads;
    using clk = std::chrono::steady_clock;
    auto ms = [](clk::time_point x, clk::time_point y) { return std::chrono::duration<double, std::milli>(y - x).count(); };
    auto t0 = clk::now();
    Encoded enc = encode(m, rgb, W, H);
    auto t1 = clk::now();
    std::vector<Box> objects, texts;
    detect(enc, o.conf, o.text_conf, o.nms, o.max_objects, o.max_text_lines, objects, texts);
    const Meta& mt = enc.meta;
    auto to1000 = [&](const Box& b, int out[4]) {
        double sx = 512.0 / mt.r / mt.W * 1000.0, sy = 512.0 / mt.r / mt.H * 1000.0;
        out[0] = (int)std::nearbyint(b.x0 * sx);
        out[1] = (int)std::nearbyint(b.y0 * sy);
        out[2] = (int)std::nearbyint(b.x1 * sx);
        out[3] = (int)std::nearbyint(b.y1 * sy);
    };

    if (!o.form.empty()) {
        std::vector<float> M;
        std::vector<uint8_t> mv;
        std::vector<int> refs;
        memory(m, enc, objects, texts, M, mv, refs);
        std::vector<std::vector<int>> ref_boxes;
        for (const auto* list : {&objects, &texts})
            for (const Box& b : *list) {
                int bb[4];
                to1000(b, bb);
                ref_boxes.push_back({bb[0], bb[1], bb[2], bb[3]});
            }
        auto t2 = clk::now();
        std::string form = fill_form(m, M, mv, refs, ref_boxes, o.language, o.form);
        auto t3 = clk::now();
        if (tm) { tm[0] = ms(t0, t1); tm[1] = ms(t1, t2); tm[2] = ms(t2, t3); tm[3] = ms(t0, t3); tm[4] = 0; }
        return "{\"form\":" + form + ",\"timing_ms\":{\"encode\":" + num(ms(t0, t1), 1) + ",\"detect\":" + num(ms(t1, t2), 1) +
               ",\"form\":" + num(ms(t2, t3), 1) + ",\"total\":" + num(ms(t0, t3), 1) + "}}";
    }

    auto wants = [&](const char* s) { return o.sections.empty() || has(o.sections, s); };
    auto field = [&](const char* f) { return o.fields.empty() || has(o.fields, f); };
    std::vector<std::string> ocr;
    for (auto& b : texts) {
        if (!wants("text")) { ocr.push_back(""); continue; }
        int bb[4];
        to1000(b, bb);
        ocr.push_back(read_text(m, rgb, W, H, bb[0] / 1000.0f * W, bb[1] / 1000.0f * H, bb[2] / 1000.0f * W, bb[3] / 1000.0f * H));
    }
    auto t2 = clk::now();
    bool need_decoder = wants("scene") || wants("objects") || wants("regions") || wants("relations") || wants("events") || wants("reasoning");
    std::vector<float> M;
    std::vector<uint8_t> mv;
    std::vector<int> refs;
    std::vector<int> ids;
    if (need_decoder) {
        memory(m, enc, objects, texts, M, mv, refs);
        ids = generate(m, M, mv, refs, o.language, o.max_tokens, o.keywords);
    }
    auto t3 = clk::now();
    std::string text = detokenize(m, ids);

    std::string trace, body = text;
    size_t a = text.find("<think>"), b = text.find("</think>");
    if (a != std::string::npos && b != std::string::npos && b > a) {
        trace = trim(text.substr(a + 7, b - a - 7));
        body = text.substr(b + 8);
    }
    std::vector<std::pair<std::string, std::string>> scene;
    std::vector<Entity> ents;
    std::vector<int> text_refs;
    std::vector<std::string> rel_json, events;
    std::stringstream ls(body);
    std::string line;
    while (std::getline(ls, line)) {
        line = trim(line);
        if (line.empty()) continue;
        if (line.rfind("S ", 0) == 0) {
            for (auto& kv : parse_attrs(line.substr(2))) scene.push_back(kv);
        } else if (line.size() > 1 && line[0] == 'E' && std::isdigit((unsigned char)line[1])) {
            Entity e;
            size_t sp = line.find(' ');
            e.n = std::atoi(line.substr(1, sp).c_str());
            std::string rest = sp == std::string::npos ? "" : line.substr(sp + 1);
            size_t br = rest.rfind('{');
            if (br != std::string::npos && !rest.empty() && rest.back() == '}') {
                e.attrs = parse_attrs(rest.substr(br + 1, rest.size() - br - 2));
                rest = trim(rest.substr(0, br));
            }
            size_t r0;
            while ((r0 = rest.find('<')) != std::string::npos) {
                size_t r1 = rest.find('>', r0);
                if (r1 == std::string::npos) break;
                std::string tag = rest.substr(r0 + 1, r1 - r0 - 1);
                if (tag.rfind("obj", 0) == 0) e.refs.push_back({'o', std::atoi(tag.c_str() + 3)});
                else if (tag.rfind("txt", 0) == 0) e.refs.push_back({'t', std::atoi(tag.c_str() + 3)});
                rest.erase(r0, r1 - r0 + 1);
            }
            for (auto& c : rest)
                if (c == '*') c = ' ';
            size_t at = rest.find('@');
            if (at != std::string::npos) { e.position = trim(rest.substr(at + 1)); rest = rest.substr(0, at); }
            rest = trim(rest);
            size_t hs = rest.rfind(" #");
            if (hs != std::string::npos) { e.count = trim(rest.substr(hs + 2)); rest = trim(rest.substr(0, hs)); }
            e.name = rest;
            ents.push_back(e);
        } else if (line[0] == 'T') {
            size_t r0 = 0;
            while ((r0 = line.find("<txt", r0)) != std::string::npos) { text_refs.push_back(std::atoi(line.c_str() + r0 + 4)); r0 += 4; }
        } else if (line.rfind("R ", 0) == 0) {
            std::stringstream rs(line.substr(2));
            int s1, s2;
            std::string pred;
            if (rs >> s1 >> pred >> s2)
                rel_json.push_back("{\"subject\":" + std::to_string(s1) + ",\"predicate\":" + esc(pred) + ",\"object\":" + std::to_string(s2) + "}");
        } else if (line.rfind("V ", 0) == 0) {
            events.push_back(trim(line.substr(2)));
        }
    }

    std::vector<float> protos;
    const bool draw = o.masks && o.mask_format != "none" && field("mask");
    auto mask_json = [&](const Box& box, const int bb[4]) -> std::string {
        std::vector<float> prob = mask_prob(m, enc, box, &protos);
        if (o.mask_format == "bitmap") return bitmap_rle(prob, bb, mt, o.mask_threshold);
        double sx = 512.0 / 128.0 / mt.r / mt.W * 1000.0, sy = 512.0 / 128.0 / mt.r / mt.H * 1000.0;
        auto ring_str = [&](const Ring& rg) {
            std::ostringstream rs;
            rs << "[";
            for (size_t i = 0; i < rg.size(); i++)
                rs << (i ? "," : "") << (int)std::nearbyint(std::clamp(rg[i].first * sx, 0.0, 1000.0)) << ","
                   << (int)std::nearbyint(std::clamp(rg[i].second * sy, 0.0, 1000.0));
            rs << "]";
            return rs.str();
        };
        std::vector<Ring> rings = outlines(prob, o.mask_threshold, o.mask_format == "largest");
        if (o.mask_format == "largest") return rings.empty() ? "[]" : ring_str(rings[0]);
        if (o.mask_format == "merged") return ring_str(merged_ring(rings));
        std::string s = "[";
        for (size_t i = 0; i < rings.size(); i++) s += (i ? "," : "") + ring_str(rings[i]);
        return s + "]";
    };
    auto common = [&](const Entity& e) {
        std::ostringstream it;
        it << "{\"id\":" << e.n << ",\"name\":" << esc(e.name);
        if (!e.count.empty() && field("count")) {
            bool digits = std::all_of(e.count.begin(), e.count.end(), ::isdigit);
            it << ",\"count\":" << (digits ? e.count : esc(e.count));
        }
        if (!e.position.empty() && field("position")) it << ",\"position\":" << esc(e.position);
        if (!e.attrs.empty() && field("description")) {
            it << ",\"description\":{";
            for (size_t i = 0; i < e.attrs.size(); i++) it << (i ? "," : "") << esc(e.attrs[i].first) << ":" << esc(e.attrs[i].second);
            it << "}";
        }
        return it.str();
    };
    // One entry per referenced instance, sharing the entity's id, name and description.
    std::ostringstream objs_js, regs_js;
    bool ofirst = true, rfirst = true;
    for (auto& e : ents) {
        std::vector<int> inst;
        for (auto& r : e.refs)
            if (r.first == 'o' && r.second >= 0 && r.second < (int)objects.size()) inst.push_back(r.second);
        if (inst.empty()) {
            regs_js << (rfirst ? "" : ",") << common(e) << "}";
            rfirst = false;
            continue;
        }
        for (int k : inst) {
            int bb[4];
            to1000(objects[k], bb);
            objs_js << (ofirst ? "" : ",") << common(e);
            if (field("bbox")) objs_js << ",\"bbox\":[" << bb[0] << "," << bb[1] << "," << bb[2] << "," << bb[3] << "]";
            if (field("score")) objs_js << ",\"score\":" << num(objects[k].score);
            if (draw) objs_js << ",\"mask\":" << mask_json(objects[k], bb);
            objs_js << "}";
            ofirst = false;
        }
    }
    std::ostringstream js;
    std::vector<std::string> parts;
    if (wants("scene")) {
        std::string s = "\"scene\":{";
        for (size_t i = 0; i < scene.size(); i++) s += (i ? "," : "") + esc(scene[i].first) + ":" + esc(scene[i].second);
        parts.push_back(s + "}");
    }
    if (wants("objects")) parts.push_back("\"objects\":[" + objs_js.str() + "]");
    if (wants("regions")) parts.push_back("\"regions\":[" + regs_js.str() + "]");
    if (wants("text")) {
        std::ostringstream ts;
        bool first = true;
        for (size_t k = 0; k < texts.size(); k++) {
            if (ocr[k].empty()) continue;
            int bb[4];
            to1000(texts[k], bb);
            bool mentioned = std::find(text_refs.begin(), text_refs.end(), (int)k) != text_refs.end();
            ts << (first ? "" : ",") << "{\"string\":" << esc(ocr[k]);
            if (field("bbox")) ts << ",\"bbox\":[" << bb[0] << "," << bb[1] << "," << bb[2] << "," << bb[3] << "]";
            if (field("score")) ts << ",\"score\":" << num(texts[k].score);
            if (need_decoder) ts << ",\"mentioned\":" << (mentioned ? "true" : "false");
            ts << "}";
            first = false;
        }
        parts.push_back("\"text\":[" + ts.str() + "]");
    }
    if (wants("relations")) {
        std::string s = "\"relations\":[";
        for (size_t i = 0; i < rel_json.size(); i++) s += (i ? "," : "") + rel_json[i];
        parts.push_back(s + "]");
    }
    if (wants("events")) {
        std::string s = "\"events\":[";
        for (size_t i = 0; i < events.size(); i++) s += (i ? "," : "") + esc(events[i]);
        parts.push_back(s + "]");
    }
    if (wants("reasoning") && o.reasoning) parts.push_back("\"reasoning\":" + esc(trace));
    auto t4 = clk::now();
    parts.push_back("\"timing_ms\":{\"encode\":" + num(ms(t0, t1), 1) + ",\"detect_read\":" + num(ms(t1, t2), 1) + ",\"decode\":" +
                    num(ms(t2, t3), 1) + ",\"assemble\":" + num(ms(t3, t4), 1) + ",\"total\":" + num(ms(t0, t4), 1) +
                    ",\"tokens\":" + std::to_string(ids.size()) + "}");
    js << "{";
    for (size_t i = 0; i < parts.size(); i++) js << (i ? "," : "") << parts[i];
    js << "}";
    if (tm) { tm[0] = ms(t0, t1); tm[1] = ms(t1, t2); tm[2] = ms(t2, t3); tm[3] = ms(t0, t4); tm[4] = (double)ids.size(); }
    return js.str();
}

// ------------------------------------------------------------------ options
static void parse_form(const JV& v, Options& o) {
    auto bad = [](const std::string& what) { throw std::runtime_error("invalid form: " + what); };
    if (v.t != JV::OBJ) bad("form must be an object of fields");
    if (v.o.empty()) bad("form has no fields");
    for (const auto& kv : v.o) {
        const std::string& name = kv.first;
        const JV& s = kv.second;
        if (name.empty() || name.find_first_of("=;|#$@\n ") != std::string::npos) bad("field name " + esc(name) + " cannot be used");
        if (s.t != JV::OBJ) bad("field " + name + " must be an object");
        FormField f;
        f.name = name;
        const JV* en = s.get("enum");
        const JV* xs = s.get("x-source");
        const JV* ty = s.get("type");
        bool done = false;
        if (en) {
            if (en->t != JV::ARR || en->a.empty()) bad("field " + name + " needs an enum list");
            for (const JV& c : en->a) {
                if (c.t == JV::STR) { f.choices.push_back(c.s); f.values.push_back(esc(c.s)); }
                else if (c.t == JV::NUM) { f.choices.push_back(c.s); f.values.push_back(c.s); }
                else bad("field " + name + " has a choice that is not a string or a number");
                if (f.choices.back().empty() || f.choices.back().find_first_of("|\n") != std::string::npos)
                    bad("field " + name + " has a choice that cannot be used");
            }
            f.kind = FormField::CHOICE;
            done = true;
        } else if (xs) {
            if (xs->t != JV::STR) bad("field " + name + " needs a string x-source");
            if (xs->s == "bbox") {
                const JV* of = s.get("x-of");
                bool ok = false;
                if (of && of->t == JV::STR)
                    for (const FormField& g : o.form) ok = ok || (g.name == of->s && g.kind == FormField::CHOICE);
                if (!ok) bad("box field " + name + " must point at an earlier choice field");
                f.kind = FormField::BOX;
                f.of = of->s;
                done = true;
            }
        }
        if (!done && ty && ty->t == JV::STR && ty->s == "integer") {
            const JV* lo = s.get("minimum");
            const JV* hi = s.get("maximum");
            if (!lo || !hi || lo->t != JV::NUM || hi->t != JV::NUM || lo->n != std::floor(lo->n) || hi->n != std::floor(hi->n) || lo->n > hi->n)
                bad("field " + name + " needs an integer minimum and maximum");
            if (hi->n - lo->n > 1000) bad("field " + name + " spans more than 1000 values");
            f.kind = FormField::COUNT;
            f.lo = (long)lo->n;
            f.hi = (long)hi->n;
            done = true;
        }
        if (!done && ty && ty->t == JV::STR && ty->s == "string") {
            const JV* ml = s.get("maxLength");
            if (ml) {
                if (ml->t != JV::NUM || ml->n < 1 || ml->n > 500) bad("field " + name + " needs a maxLength from 1 to 500");
                f.max_length = (int)ml->n;
            }
            f.kind = FormField::TEXT;
            done = true;
        }
        if (!done) bad("field " + name + " is not a choice, count, box or text field");
        o.form.push_back(f);
    }
}

static void parse_options(const Model& m, const char* s, Options& o) {
    if (!s || !*s) return;
    JReader jr{s, s + std::strlen(s)};
    JV root = jr.value();
    jr.ws();
    if (jr.p != jr.e) jr.fail("trailing characters after the JSON object");
    if (root.t != JV::OBJ) jr.fail("options must be a JSON object");
    static const char* SECTIONS[] = {"scene", "objects", "regions", "text", "relations", "events", "reasoning"};
    static const char* FIELDS[] = {"bbox", "score", "mask", "position", "description", "count"};
    for (const auto& kv : root.o) {
        const std::string& k = kv.first;
        const JV& v = kv.second;
        auto number = [&](double lo, double hi) {
            if (v.t != JV::NUM || !(v.n >= lo && v.n <= hi)) jr.fail(k + " must be a number from " + num(lo, 2) + " to " + num(hi, 2));
            return v.n;
        };
        auto boolean = [&]() {
            if (v.t != JV::BOOL) jr.fail(k + " must be true or false");
            return v.b;
        };
        auto names = [&](const char* const* allowed, size_t n, std::vector<std::string>& out) {
            if (v.t != JV::ARR) jr.fail(k + " must be a list of names");
            for (const JV& x : v.a) {
                bool ok = x.t == JV::STR && std::find(allowed, allowed + n, x.s) != allowed + n;
                if (!ok) jr.fail(k + " has an unknown name " + (x.t == JV::STR ? x.s : std::string("(not a string)")));
                out.push_back(x.s);
            }
        };
        if (k == "language") {
            if (v.t != JV::STR || std::find(m.languages.begin(), m.languages.end(), v.s) == m.languages.end()) {
                std::string list;
                for (const auto& l : m.languages) list += (list.empty() ? "" : ", ") + l;
                jr.fail("unsupported language " + (v.t == JV::STR ? v.s : std::string("(not a string)")) + "; choose one of " + list);
            }
            o.language = v.s;
        } else if (k == "keywords") {
            std::vector<std::string> ks;
            if (v.t == JV::STR) {
                std::stringstream ss(v.s);
                std::string line;
                while (std::getline(ss, line)) ks.push_back(line);
            } else if (v.t == JV::ARR) {
                for (const JV& x : v.a) {
                    if (x.t != JV::STR) jr.fail("keywords must be strings");
                    ks.push_back(x.s);
                }
            } else {
                jr.fail("keywords must be a list of strings or one newline-separated string");
            }
            for (auto& s : ks) {
                s = trim(s);
                if (s.empty()) continue;
                if (s.size() > 100) jr.fail("a keyword is longer than 100 bytes");
                o.keywords.push_back(s);
            }
            if (o.keywords.size() > 64) jr.fail("more than 64 keywords");
        } else if (k == "conf") o.conf = (float)number(0, 1);
        else if (k == "text_conf") o.text_conf = (float)number(0, 1);
        else if (k == "mask_threshold") o.mask_threshold = (float)number(0, 1);
        else if (k == "max_objects") o.max_objects = (int)number(0, 64);
        else if (k == "max_text_lines") o.max_text_lines = (int)number(0, 32);
        else if (k == "max_tokens") o.max_tokens = (int)number(1, 1536);
        else if (k == "masks") o.masks = boolean();
        else if (k == "nms") o.nms = boolean();
        else if (k == "reasoning") o.reasoning = boolean();
        else if (k == "mask_format") {
            if (v.t != JV::STR || (v.s != "largest" && v.s != "polygons" && v.s != "merged" && v.s != "bitmap" && v.s != "none"))
                jr.fail("mask_format must be largest, polygons, merged, bitmap or none");
            o.mask_format = v.s;
        } else if (k == "sections") names(SECTIONS, 7, o.sections);
        else if (k == "fields") names(FIELDS, 6, o.fields);
        else if (k == "form") parse_form(v, o);
    }
}

}  // namespace oasis

// ------------------------------------------------------------------ C API
static thread_local std::string g_error;

struct oasis_model {
    oasis::Model m;
};

extern "C" {

oasis_model* oasis_load(const char* path) {
    const char* ex = std::getenv("OASIS_EXACT");
    oasis::g_exact_gemv = ex && ex[0] == '1';
    auto* h = new oasis_model();
    std::string err;
    try {
        if (!h->m.load(path, err)) {
            g_error = err;
            delete h;
            return nullptr;
        }
    } catch (const std::exception& e) {
        g_error = e.what();
        delete h;
        return nullptr;
    }
    return h;
}

void oasis_free(oasis_model* h) { delete h; }

int oasis_describe(oasis_model* h, const unsigned char* rgb, int width, int height, const char* options_json, char* out, int capacity) {
    if (!h || !rgb || width <= 0 || height <= 0) { g_error = "invalid arguments"; return -1; }
    try {
        oasis::Options o;
        oasis::parse_options(h->m, options_json, o);
        std::string js = oasis::describe_json(h->m, rgb, width, height, o, nullptr);
        if ((int)js.size() + 1 > capacity) { g_error = "output buffer too small"; return -(int)js.size() - 1; }
        std::memcpy(out, js.c_str(), js.size() + 1);
        return (int)js.size();
    } catch (const std::exception& e) {
        g_error = e.what();
        return -1;
    }
}

const char* oasis_last_error(void) { return g_error.c_str(); }

}  // extern "C"