// bp.c - runs Basic Pitch from a plan (gen.mjs): a list of ops on 4-D NHWC tensors.

#include "bp.h"
#include <float.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define MIN(a, b) ((a) < (b) ? (a) : (b))
#define MAX(a, b) ((a) > (b) ? (a) : (b))

#define PLAN_MAGIC 0x42504c4f
#define FRAMES     172                       // rows per window
#define EDGE       15                        // rows dropped at each end
#define OVERLAP    (2 * EDGE * BP_HOP)       // samples two windows share
#define STEP       (BP_WINDOW - OVERLAP)
#define MAX_OC     64                        // most conv output channels

enum { BUFFER, WEIGHT, ALIAS, INPUT };    // tensor kinds
enum { CONV = 1, PAD, MIRROR_PAD, TRANSPOSE, SLICE, CONCAT, UNARY, BINARY, REDUCE, FREE, HCONV };

typedef struct {
    int kind, ref;   // ref: weight offset (WEIGHT) or the aliased tensor (ALIAS)
    int d[4];
    float *data;
} tensor;

struct bp_model {
    void *owned[2];   // plan and weights read by bp_load_files
    tensor *t;
    int nt, nops, frames, onsets, input, contours;
    const int *ops;
    struct { const int *op; float *taps; } merged[4];   // HCONV filters
};

static int size4(const int *d) { return d[0] * d[1] * d[2] * d[3]; }
static float *data(const bp_model *m, const tensor *t) { return t->kind == ALIAS ? data(m, &m->t[t->ref]) : t->data; }
static float sigmoid(float x) { return 1 / (1 + expf(-x)); }
static float activate(float x, int act) { return act == 1 ? (x > 0 ? x : 0) : act == 2 ? sigmoid(x) : x; }

// plan: magic, tensor count, op count, the frames/onsets/input/contours tensors, 6 words a tensor
// (kind, ref, shape), then the ops
bp_model *bp_load(const int *plan, int words, const float *weights) {
    if (words < 7 || plan[0] != PLAN_MAGIC || plan[1] <= 0 || 7 + 6 * plan[1] > words) return NULL;
    bp_model *m = calloc(1, sizeof *m);
    if (!m || !(m->t = calloc(plan[1], sizeof(tensor)))) { free(m); return NULL; }
    m->nt = plan[1]; m->nops = plan[2]; m->frames = plan[3]; m->onsets = plan[4]; m->input = plan[5]; m->contours = plan[6];
    const int *p = plan + 7;
    for (int i = 0; i < m->nt; i++, p += 6) {
        m->t[i] = (tensor){ p[0], p[1], { p[2], p[3], p[4], p[5] }, NULL };
        if (m->t[i].kind == WEIGHT) m->t[i].data = (float *)weights + m->t[i].ref;
    }
    m->ops = p;
    return m;
}

static void *read_file(const char *path, long *size) {
    FILE *f = fopen(path, "rb");
    if (!f) return NULL;
    void *data = NULL;
    if (fseek(f, 0, SEEK_END) == 0 && (*size = ftell(f)) > 0 && fseek(f, 0, SEEK_SET) == 0 && (data = malloc(*size)) && fread(data, 1, *size, f) != (size_t)*size) {
        free(data);
        data = NULL;
    }
    fclose(f);
    return data;
}

bp_model *bp_load_files(const char *plan_path, const char *weights_path) {
    long plan_size, weights_size;
    int *plan = read_file(plan_path, &plan_size);
    float *weights = read_file(weights_path, &weights_size);
    bp_model *m = plan && weights ? bp_load(plan, plan_size / 4, weights) : NULL;
    if (!m) { free(plan); free(weights); return NULL; }
    m->owned[0] = plan;
    m->owned[1] = weights;
    return m;
}

void bp_unload(bp_model *m) {
    if (!m) return;
    free(m->owned[0]);
    free(m->owned[1]);
    for (int i = 0; i < 4; i++) free(m->merged[i].taps);
    free(m->t);
    free(m);
}

// conv params: KH KW stride_h stride_w dilation_h dilation_w pad_top pad_left activation
// (0, relu, sigmoid) has_bias. NHWC input, HWIO filter.

typedef float v4 __attribute__((vector_size(16)));   // wasm SIMD with -msimd128
static inline v4 load4(const float *p) { v4 v; memcpy(&v, p, sizeof v); return v; }
static inline v4 splat(float x) { return (v4){ x, x, x, x }; }

// one output pixel, inlined per OC so the channel loops unroll
static inline __attribute__((always_inline)) void conv_pixel(const int OC, const float *restrict in, const float *restrict wt, const float *restrict bias,
                                                             float *restrict out, int oh, int ow, const int *p, int H, int W, int C, int OW) {
    const int KH = p[0], KW = p[1], sh = p[2], sw = p[3], dh = p[4], dw = p[5], top = p[6], left = p[7];
    float sum[MAX_OC];
    v4 vsum[MAX_OC / 4];
    for (int oc = 0; oc < OC; oc++) sum[oc] = bias ? bias[oc] : 0;
    for (int g = 0; g < OC / 4; g++) vsum[g] = splat(0);
    int kw0 = 0, kw1 = KW;   // taps inside the input
    while (kw0 < KW && ow * sw - left + kw0 * dw < 0) kw0++;
    while (kw1 > kw0 && ow * sw - left + (kw1 - 1) * dw >= W) kw1--;
    for (int kh = 0; kh < KH; kh++) {
        const int ih = oh * sh - top + kh * dh;
        if (ih < 0 || ih >= H) continue;
        const float *row = in + (size_t)ih * W * C, *wrow = wt + (size_t)kh * KW * C * OC;
        if (OC == 1 && C == 1 && dw == 1) {
            const float *xs = row + (ow * sw - left + kw0), *ws = wrow + kw0;
            float s = 0;
            for (int k = 0; k < kw1 - kw0; k++) s += xs[k] * ws[k];
            sum[0] += s;
        } else if (OC == 1) {
            for (int kw = kw0; kw < kw1; kw++) {
                const float *xs = row + (size_t)(ow * sw - left + kw * dw) * C, *ws = wrow + (size_t)kw * C;
                float s = 0;
                for (int c = 0; c < C; c++) s += xs[c] * ws[c];
                sum[0] += s;
            }
        } else if (OC % 4 == 0) {
            for (int kw = kw0; kw < kw1; kw++) {
                const float *xs = row + (size_t)(ow * sw - left + kw * dw) * C, *ws = wrow + (size_t)kw * C * OC;
                for (int c = 0; c < C; c++) {
                    const v4 v = splat(xs[c]);
                    for (int g = 0; g < OC / 4; g++) vsum[g] += v * load4(ws + (size_t)c * OC + 4 * g);
                }
            }
        } else {
            for (int kw = kw0; kw < kw1; kw++) {
                const float *xs = row + (size_t)(ow * sw - left + kw * dw) * C, *ws = wrow + (size_t)kw * C * OC;
                for (int c = 0; c < C; c++) {
                    const float v = xs[c], *wc = ws + (size_t)c * OC;
                    for (int oc = 0; oc < OC; oc++) sum[oc] += v * wc[oc];
                }
            }
        }
    }
    if (OC % 4 == 0) for (int g = 0; g < OC / 4; g++) for (int k = 0; k < 4; k++) sum[4 * g + k] += vsum[g][k];
    float *o = out + ((size_t)oh * OW + ow) * OC;
    for (int oc = 0; oc < OC; oc++) o[oc] = activate(sum[oc], p[8]);
}

// four single-channel outputs at once, each summed in conv_pixel's order (same results)
static void conv1_x4(const float *in, const float *wt, const float *bias, float *out, int oh, int ow, const int *p, int H, int W, int C, int OW) {
    const int KH = p[0], KW = p[1], sh = p[2], sw = p[3], dh = p[4], dw = p[5], top = p[6], left = p[7];
    v4 sum = splat(bias ? bias[0] : 0);
    for (int kh = 0; kh < KH; kh++) {
        const int ih = oh * sh - top + kh * dh;
        if (ih < 0 || ih >= H) continue;
        const float *x0 = in + ((size_t)ih * W + ow * sw - left) * C, *x1 = x0 + sw * C, *x2 = x1 + sw * C, *x3 = x2 + sw * C;
        const float *w = wt + (size_t)kh * KW * C;
        if (C == 1 && dw == 1) {
            v4 s = splat(0);
            for (int k = 0; k < KW; k++) s += (v4){ x0[k], x1[k], x2[k], x3[k] } * splat(w[k]);
            sum += s;
        } else {
            for (int kw = 0; kw < KW; kw++) {
                const size_t o = (size_t)kw * dw * C;
                v4 s = splat(0);
                for (int c = 0; c < C; c++) s += (v4){ x0[o + c], x1[o + c], x2[o + c], x3[o + c] } * splat(w[kw * C + c]);
                sum += s;
            }
        }
    }
    for (int i = 0; i < 4; i++) out[(size_t)oh * OW + ow + i] = activate(sum[i], p[8]);
}

static void conv1(const float *in, const float *wt, const float *bias, float *out, int H, int W, int C, int OH, int OW, const int *p) {
    const int KW = p[1], sw = p[3], dw = p[5], left = p[7];
    for (int oh = 0; oh < OH; oh++)
        for (int ow = 0; ow < OW; ow++)
            if (ow + 3 < OW && ow * sw - left >= 0 && (ow + 3) * sw - left + (KW - 1) * dw < W) conv1_x4(in, wt, bias, out, oh, ow, p, H, W, C, OW), ow += 3;
            else conv_pixel(1, in, wt, bias, out, oh, ow, p, H, W, C, OW);
}

#define CONV_LOOP(N) for (int oh = 0; oh < OH; oh++) for (int ow = 0; ow < OW; ow++) conv_pixel(N, in, wt, bias, out, oh, ow, p, H, W, C, OW)
static void conv(const float *in, const int *xd, const float *wt, const float *bias, float *out, const int *yd, const int *p) {
    const int H = xd[1], W = xd[2], C = xd[3], OH = yd[1], OW = yd[2], OC = yd[3];
    switch (OC) {
    case 1: conv1(in, wt, bias, out, H, W, C, OH, OW, p); break;
    case 8: CONV_LOOP(8); break;
    case 32: CONV_LOOP(32); break;
    case 36: CONV_LOOP(36); break;
    default: CONV_LOOP(OC); break;
    }
}

// HCONV: harmonic stacking (H shifted copies of x [1, T, F, 1]) and the conv after it, with the
// taps that read the same bin of x merged into one filter. Params: the conv's, H, the H shifts.

static const float *merged_taps(bp_model *m, const int *p, const float *wt, int OC, int *jmin, int *ntaps) {
    const int KH = p[0], KW = p[1], left = p[7], H = p[10], *shift = p + 11;
    int lo = shift[0], hi = shift[0];
    for (int h = 1; h < H; h++) { lo = MIN(lo, shift[h]); hi = MAX(hi, shift[h]); }
    *jmin = lo - left;
    *ntaps = hi - lo + KW;
    int slot = 0;
    while (slot < 3 && m->merged[slot].op && m->merged[slot].op != p) slot++;
    if (m->merged[slot].op == p) return m->merged[slot].taps;
    float *taps = calloc((size_t)KH * *ntaps * OC, sizeof(float));
    for (int kh = 0; kh < KH; kh++)
        for (int kw = 0; kw < KW; kw++)
            for (int h = 0; h < H; h++) {
                const float *w = wt + (((size_t)kh * KW + kw) * H + h) * OC;
                float *to = taps + ((size_t)kh * *ntaps + (kw - left + shift[h] - *jmin)) * OC;
                for (int oc = 0; oc < OC; oc++) to[oc] += w[oc];
            }
    free(m->merged[slot].taps);
    m->merged[slot].op = p;
    m->merged[slot].taps = taps;
    return taps;
}

// four outputs at once where the merged filter fits, as conv1_x4
static inline __attribute__((always_inline)) void hconv_x4(const int OC, const float *restrict x, const float *restrict taps, const float *restrict bias,
                                                           float *restrict out, int oh, int ow, const int *p, int T, int F, int OW, int jmin, int ntaps) {
    const int KH = p[0], top = p[6];
    float sum[4][MAX_OC];
    for (int i = 0; i < 4; i++) for (int oc = 0; oc < OC; oc++) sum[i][oc] = bias ? bias[oc] : 0;
    for (int kh = 0; kh < KH; kh++) {
        const int t = oh - top + kh;
        if (t < 0 || t >= T) continue;
        const float *row = x + (size_t)t * F + ow + jmin;
        v4 vsum[4][MAX_OC / 4];
        for (int i = 0; i < 4; i++) for (int g = 0; g < OC / 4; g++) vsum[i][g] = splat(0);
        for (int j = 0; j < ntaps; j++) {
            const float *mj = taps + ((size_t)kh * ntaps + j) * OC;
            for (int i = 0; i < 4; i++) {
                const v4 v = splat(row[i + j]);
                for (int g = 0; g < OC / 4; g++) vsum[i][g] += v * load4(mj + 4 * g);
            }
        }
        for (int i = 0; i < 4; i++) for (int g = 0; g < OC / 4; g++) for (int k = 0; k < 4; k++) sum[i][4 * g + k] += vsum[i][g][k];
    }
    for (int i = 0; i < 4; i++) {
        float *o = out + ((size_t)oh * OW + ow + i) * OC;
        for (int oc = 0; oc < OC; oc++) o[oc] = activate(sum[i][oc], p[8]);
    }
}

static inline __attribute__((always_inline)) void hconv_row(const int OC, const float *restrict x, const float *restrict wt, const float *restrict taps,
                                                            const float *restrict bias, float *restrict out, int oh, const int *p, int T, int F, int OW, int jmin, int ntaps) {
    const int KH = p[0], KW = p[1], top = p[6], left = p[7], H = p[10], *shift = p + 11;
    for (int ow = 0; ow < OW; ow++) {
        if (OC == 8 && ow + 3 < OW && ow - left >= 0 && ow + 3 - left + KW <= OW && ow >= -jmin && ow + 3 + jmin + ntaps <= F) {
            hconv_x4(OC, x, taps, bias, out, oh, ow, p, T, F, OW, jmin, ntaps);
            ow += 3;
            continue;
        }
        float sum[MAX_OC];
        for (int oc = 0; oc < OC; oc++) sum[oc] = bias ? bias[oc] : 0;
        const int inside = ow - left >= 0 && ow - left + KW <= OW;
        for (int kh = 0; kh < KH; kh++) {
            const int t = oh - top + kh;
            if (t < 0 || t >= T) continue;
            const float *row = x + (size_t)t * F;
            if (inside) {
                const int j0 = MAX(jmin, -ow), j1 = MIN(jmin + ntaps - 1, F - 1 - ow);
                if (OC % 4 == 0) {
                    v4 vsum[MAX_OC / 4];
                    for (int g = 0; g < OC / 4; g++) vsum[g] = splat(0);
                    for (int j = j0; j <= j1; j++) {
                        const float *mj = taps + ((size_t)kh * ntaps + (j - jmin)) * OC;
                        const v4 v = splat(row[ow + j]);
                        for (int g = 0; g < OC / 4; g++) vsum[g] += v * load4(mj + 4 * g);
                    }
                    for (int g = 0; g < OC / 4; g++) for (int k = 0; k < 4; k++) sum[4 * g + k] += vsum[g][k];
                } else {
                    for (int j = j0; j <= j1; j++) {
                        const float v = row[ow + j], *mj = taps + ((size_t)kh * ntaps + (j - jmin)) * OC;
                        for (int oc = 0; oc < OC; oc++) sum[oc] += v * mj[oc];
                    }
                }
            } else {
                for (int kw = 0; kw < KW; kw++) {
                    const int g = ow - left + kw;
                    if (g < 0 || g >= OW) continue;
                    for (int h = 0; h < H; h++) {
                        const int i = g + shift[h];
                        if (i < 0 || i >= F) continue;
                        const float v = row[i], *w = wt + (((size_t)kh * KW + kw) * H + h) * OC;
                        for (int oc = 0; oc < OC; oc++) sum[oc] += v * w[oc];
                    }
                }
            }
        }
        float *o = out + ((size_t)oh * OW + ow) * OC;
        for (int oc = 0; oc < OC; oc++) o[oc] = activate(sum[oc], p[8]);
    }
}

static void hconv(bp_model *m, const float *x, const int *xd, const float *wt, const float *bias, float *out, const int *yd, const int *p) {
    const int T = xd[1], F = xd[2], OH = yd[1], OW = yd[2], OC = yd[3];
    int jmin, ntaps;
    const float *taps = merged_taps(m, p, wt, OC, &jmin, &ntaps);
    for (int oh = 0; oh < OH; oh++)
        if (OC == 8) hconv_row(8, x, wt, taps, bias, out, oh, p, T, F, OW, jmin, ntaps);
        else hconv_row(OC, x, wt, taps, bias, out, oh, p, T, F, OW, jmin, ntaps);
}

#define FOR4(d) for (int i0 = 0; i0 < (d)[0]; i0++) for (int i1 = 0; i1 < (d)[1]; i1++) for (int i2 = 0; i2 < (d)[2]; i2++) for (int i3 = 0; i3 < (d)[3]; i3++)
#define AT(d, a, b, c, e) ((((size_t)(a) * (d)[1] + (b)) * (d)[2] + (c)) * (d)[3] + (e))

static int reflect(int i, int n) { return i < 0 ? -i : i >= n ? 2 * n - 2 - i : i; }

static void run_op(bp_model *m, int code, tensor *y, tensor **x, int nin, const int *p) {
    const int *xd = nin ? x[0]->d : NULL, *yd = y->d;
    const float *a = nin ? data(m, x[0]) : NULL, *bias = nin > 2 ? data(m, x[2]) : NULL;
    float *out = y->data;
    switch (code) {
    case CONV: conv(a, xd, data(m, x[1]), bias, out, yd, p); break;
    case HCONV: hconv(m, a, xd, data(m, x[1]), bias, out, yd, p); break;
    case PAD: case MIRROR_PAD:   // p: padding before each axis
        FOR4(yd) {
            int s[4] = { i0 - p[0], i1 - p[1], i2 - p[2], i3 - p[3] }, inside = 1;
            for (int k = 0; k < 4; k++) {
                if (code == MIRROR_PAD) s[k] = reflect(s[k], xd[k]);
                else if (s[k] < 0 || s[k] >= xd[k]) inside = 0;
            }
            out[AT(yd, i0, i1, i2, i3)] = inside ? a[AT(xd, s[0], s[1], s[2], s[3])] : 0;
        }
        break;
    case TRANSPOSE:   // p: the permutation
        FOR4(yd) {
            int i[4] = { i0, i1, i2, i3 }, j[4];
            for (int k = 0; k < 4; k++) j[p[k]] = i[k];
            out[AT(yd, i0, i1, i2, i3)] = a[AT(xd, j[0], j[1], j[2], j[3])];
        }
        break;
    case SLICE:       // p: begin[4], stride[4]
        FOR4(yd) out[AT(yd, i0, i1, i2, i3)] = a[AT(xd, p[0] + i0 * p[4], p[1] + i1 * p[5], p[2] + i2 * p[6], p[3] + i3 * p[7])];
        break;
    case CONCAT: {    // p: the axis
        size_t outer = 1, inner = 1, at = 0;
        for (int k = 0; k < p[0]; k++) outer *= yd[k];
        for (int k = p[0] + 1; k < 4; k++) inner *= yd[k];
        for (size_t o = 0; o < outer; o++)
            for (int n = 0; n < nin; n++) {
                const size_t len = (size_t)x[n]->d[p[0]] * inner;
                memcpy(out + at, data(m, x[n]) + o * len, len * sizeof(float));
                at += len;
            }
        break;
    }
    case UNARY: {     // p: neg, square, sqrt, log, sigmoid, relu
        const int n = size4(yd);
        for (int i = 0; i < n; i++) {
            const float v = a[i];
            out[i] = p[0] == 1 ? -v : p[0] == 2 ? v * v : p[0] == 3 ? sqrtf(v) : p[0] == 4 ? logf(v) : p[0] == 5 ? sigmoid(v) : (v > 0 ? v : 0);
        }
        break;
    }
    case BINARY: {    // p: mul, add, sub, div (x / 0 = 0), broadcasting
        const float *b = data(m, x[1]);
        const int *bd = x[1]->d;
        FOR4(yd) {
            const float u = a[AT(xd, xd[0] > 1 ? i0 : 0, xd[1] > 1 ? i1 : 0, xd[2] > 1 ? i2 : 0, xd[3] > 1 ? i3 : 0)];
            const float v = b[AT(bd, bd[0] > 1 ? i0 : 0, bd[1] > 1 ? i1 : 0, bd[2] > 1 ? i2 : 0, bd[3] > 1 ? i3 : 0)];
            out[AT(yd, i0, i1, i2, i3)] = p[0] == 1 ? u * v : p[0] == 2 ? u + v : p[0] == 3 ? u - v : (v != 0 ? u / v : 0);
        }
        break;
    }
    case REDUCE: {    // p: sum, min or max, mask of the reduced axes
        const int n = size4(yd), mask = p[1];
        for (int i = 0; i < n; i++) out[i] = p[0] == 1 ? 0 : p[0] == 2 ? FLT_MAX : -FLT_MAX;
        FOR4(xd) {
            float *o = out + AT(yd, mask & 1 ? 0 : i0, mask & 2 ? 0 : i1, mask & 4 ? 0 : i2, mask & 8 ? 0 : i3);
            const float v = a[AT(xd, i0, i1, i2, i3)];
            *o = p[0] == 1 ? *o + v : p[0] == 2 ? fminf(*o, v) : fmaxf(*o, v);
        }
        break;
    }
    }
}

// an op: code, output, input count, inputs, param count, params
int bp_run(bp_model *m, const float *window, float *frames, float *onsets, float *contours) {
    if (!m) return -1;
    m->t[m->input].data = (float *)window;
    int err = 0;
    const int *p = m->ops;
    for (int k = 0; k < m->nops && !err; k++) {
        const int code = p[0], nin = p[2];
        tensor *y = &m->t[p[1]], *x[16];
        if (nin > 16) { err = -2; break; }
        for (int i = 0; i < nin; i++) x[i] = &m->t[p[3 + i]];
        if (code == FREE) {
            free(y->data);
            y->data = NULL;
        } else if (!(y->data = malloc((size_t)size4(y->d) * sizeof(float)))) {
            err = -3;
        } else {
            run_op(m, code, y, x, nin, p + 4 + nin);
        }
        p += 4 + nin + p[3 + nin];
    }
    if (!err) {
        memcpy(frames, data(m, &m->t[m->frames]), FRAMES * BP_NOTES * sizeof(float));
        memcpy(onsets, data(m, &m->t[m->onsets]), FRAMES * BP_NOTES * sizeof(float));
        if (contours) memcpy(contours, data(m, &m->t[m->contours]), FRAMES * BP_BINS * sizeof(float));
    }
    for (int i = 0; i < m->nt; i++)
        if (m->t[i].kind == BUFFER) { free(m->t[i].data); m->t[i].data = NULL; }
    return err;
}

// as Basic Pitch: half an overlap of silence first, then a window every STEP samples
int bp_window_count(int samples) { return (samples + OVERLAP / 2 + STEP - 1) / STEP; }
int bp_frame_count(int samples) { return (int)((long long)samples * (BP_SAMPLE_RATE / BP_HOP) / BP_SAMPLE_RATE); }

int bp_window(bp_model *m, const float *audio, int samples, int index, float *frames, float *onsets, float *contours) {
    if (!m) return -1;
    float *window = calloc(BP_WINDOW, sizeof(float));
    float *f = malloc(FRAMES * BP_NOTES * sizeof(float)), *o = malloc(FRAMES * BP_NOTES * sizeof(float));
    float *c = contours ? malloc(FRAMES * BP_BINS * sizeof(float)) : NULL;
    int err = -3;
    if (window && f && o && (c || !contours)) {
        const int start = index * STEP - OVERLAP / 2;
        for (int k = 0; k < BP_WINDOW; k++)
            if (start + k >= 0 && start + k < samples) window[k] = audio[start + k];
        err = bp_run(m, window, f, o, c);
    }
    if (!err) {
        memcpy(frames, f + EDGE * BP_NOTES, BP_WINDOW_ROWS * BP_NOTES * sizeof(float));
        memcpy(onsets, o + EDGE * BP_NOTES, BP_WINDOW_ROWS * BP_NOTES * sizeof(float));
        if (c) memcpy(contours, c + EDGE * BP_BINS, BP_WINDOW_ROWS * BP_BINS * sizeof(float));
    }
    free(window); free(f); free(o); free(c);
    return err ? err : BP_WINDOW_ROWS;
}

int bp_transcribe(bp_model *m, const float *audio, int samples, float *frames, float *onsets, float *contours) {
    if (!m) return -1;
    const int rows = bp_frame_count(samples);
    float *f = malloc(BP_WINDOW_ROWS * BP_NOTES * sizeof(float)), *o = malloc(BP_WINDOW_ROWS * BP_NOTES * sizeof(float));
    float *c = contours ? malloc(BP_WINDOW_ROWS * BP_BINS * sizeof(float)) : NULL;
    int result = f && o && (c || !contours) ? rows : -3;
    for (int w = 0, at = 0; at < rows && result >= 0; w++, at += BP_WINDOW_ROWS) {
        const int err = bp_window(m, audio, samples, w, f, o, c);
        if (err < 0) { result = err; break; }
        const int n = MIN(BP_WINDOW_ROWS, rows - at);
        memcpy(frames + (size_t)at * BP_NOTES, f, n * BP_NOTES * sizeof(float));
        memcpy(onsets + (size_t)at * BP_NOTES, o, n * BP_NOTES * sizeof(float));
        if (c) memcpy(contours + (size_t)at * BP_BINS, c, n * BP_BINS * sizeof(float));
    }
    free(f); free(o); free(c);
    return result;
}