/* Microbenchmark: old (full-qsort) vs new (quickselect partial-select) DSA top-keep.
 *
 * This is NOT a unit test -- test_dsa_select.c proves correctness. This measures the
 * headline claim of #356: that replacing the O(nk log nk) qsort over all nk context
 * scores with an O(nk) partial_select_desc is materially faster per call, which is
 * the win the issue was opened for -- and that the win GROWS with context length
 * (because quickselect is linear average, qsort is n-log-n).
 *
 * It re-implements the OLD top-keep inline (qsort + threshold + scans) on a private
 * buffer so the A/B runs in one process, same inputs, same warm caches -- a controlled
 * comparison. It calls the REAL (new) partial_select_desc via the include-glm.c
 * pattern, replicating the production threshold derivation + scans.
 *
 * Methodology (chosen to be honest, not to flatter the change):
 *   - keep = 2048 (the real GLM-5.2 index_topk), nk swept across context lengths from
 *     the 2049 activation boundary up to 65536 (a long conversation).
 *   - Three score shapes: (a) realistic peaked -- a few hot keys, long tail, the shape
 *     real DSA attention scores take; (b) uniform random -- no structure; (c) a plateau
 *     of ties, to exercise the boundary-membership path.
 *   - Each (shape, nk) is timed over N_REPEAT=2000 iterations, with the scores frozen
 *     so both algorithms do IDENTICAL work. We report median ns/call and the new/old
 *     ratio. A warmup pass primes caches before timing.
 *
 * Run:  make tests/bench_dsa_select && ./tests/bench_dsa_select  (not in TEST_BINS)
 */
#define main coli_glm_main_unused
#include "../glm.c"
#undef main

#include <math.h>
#include <stdint.h>

/* ---- the OLD algorithm, verbatim from dev before #356, on a private buffer ---- */
static int cmp_pdesc_old(const void *a, const void *b){
    float x=*(const float*)a, y=*(const float*)b; return x<y?1:x>y?-1:0; }
static void keep_old(const float *isc, int nk, int keep, int *dst, int *nd_out){
    float *tmp=malloc((size_t)nk*sizeof(float));
    memcpy(tmp,isc,(size_t)nk*sizeof(float));
    qsort(tmp,(size_t)nk,sizeof(float),cmp_pdesc_old);
    float thr=tmp[keep-1]; int nd=0;
    for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
    for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
    free(tmp); *nd_out=nd;
}

/* ---- timing: median of N_REPEAT runs in ns/call, sorted ascending ---- */
#define N_REPEAT 2000
static double bench_ns(void (*fn)(const float*,int,int,int*,int*),
                       const float *isc, int nk, int keep){
    static double ts[N_REPEAT]; int *dst=malloc((size_t)nk*sizeof(int)); int nd;
    for(int r=0;r<N_REPEAT;r++){
        double t0=now_s();
        fn(isc,nk,keep,dst,&nd);
        ts[r]=(now_s()-t0)*1e9;
    }
    for(int a=1;a<N_REPEAT;a++){ double k=ts[a]; int b=a-1;
        while(b>=0 && ts[b]>k){ ts[b+1]=ts[b]; b--; } ts[b+1]=k; }
    free(dst);
    return ts[N_REPEAT/2];
}

/* the NEW algorithm calls the real partial_select_desc + replicates the production
 * threshold derivation and position scans. */
static void keep_new(const float *isc, int nk, int keep, int *dst, int *nd_out){
    float *tmp=malloc((size_t)nk*sizeof(float));
    memcpy(tmp,isc,(size_t)nk*sizeof(float));
    partial_select_desc(tmp,nk,keep);
    float thr=tmp[0]; for(int t=1;t<keep;t++) if(tmp[t]<thr) thr=tmp[t];
    int nd=0;
    for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
    for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
    free(tmp); *nd_out=nd;
}

/* deterministic score fill for three shapes */
static uint32_t brng = 0xA5A5A5A5u;
static double brand(void){ brng ^= brng << 13; brng ^= brng >> 17; brng ^= brng << 5;
    return (double)(brng >> 8) * (1.0 / 16777216.0); }
static void fill_realistic(float *isc, int nk){   /* few hot, long distinct tail */
    for(int i=0;i<nk;i++) isc[i]=(float)(-1.0 - brand()*4.0);
    isc[0]=3.f; isc[nk/50<nk?nk/50:nk-1]=1.f; isc[nk/200<nk?nk/200:nk-1]=0.5f;
}
static void fill_uniform(float *isc, int nk){     /* no structure */
    for(int i=0;i<nk;i++) isc[i]=(float)(brand()*1000.0);
}
static void fill_plateau(float *isc, int nk){     /* tie blocks -> boundary path */
    for(int i=0;i<nk;i++) isc[i]=(float)(-(double)(i/7));
}

int main(void){
    int keep = 2048;   /* GLM-5.2 index_topk */
    int nks[] = {2049, 4096, 8192, 16384, 32768, 65536};
    float *isc = malloc((size_t)65536*sizeof(float));
    int *dst   = malloc((size_t)65536*sizeof(int)); int nd;

    struct { const char *name; void (*fill)(float*,int); } shapes[] = {
        { "realistic", fill_realistic },
        { "uniform",   fill_uniform   },
        { "plateau",   fill_plateau   },
    };

    printf("bench_dsa_select: DSA top-keep, old (qsort) vs new (partial-select)  keep=%d\n", keep);
    printf("%-12s %7s %14s %14s %9s\n", "shape", "nk", "old ns/call", "new ns/call", "speedup");
    printf("------------------------------------------------------------------------\n");

    for(size_t sh=0; sh<sizeof(shapes)/sizeof(shapes[0]); sh++){
        for(size_t ni=0; ni<sizeof(nks)/sizeof(nks[0]); ni++){
            int nk=nks[ni];
            shapes[sh].fill(isc,nk);
            /* warmup both paths so caches/branch predictors are primed */
            for(int w=0; w<50; w++){ keep_old(isc,nk,keep,dst,&nd); keep_new(isc,nk,keep,dst,&nd); }
            /* sanity: both must keep exactly `keep` (correctness is test_dsa_select's
             * job, but a count divergence here would make the timing meaningless) */
            keep_old(isc,nk,keep,dst,&nd); int na=nd;
            keep_new(isc,nk,keep,dst,&nd); int nb=nd;
            if(na!=keep || nb!=keep){
                printf("%-12s %7d   (BAD COUNTS: old=%d new=%d, skipped)\n",
                       shapes[sh].name, nk, na, nb);
                continue;
            }
            double t_old=bench_ns(keep_old,isc,nk,keep);
            double t_new=bench_ns(keep_new,isc,nk,keep);
            printf("%-12s %7d %14.0f %14.0f %8.2fx\n",
                   shapes[sh].name, nk, t_old, t_new, t_old/t_new);
        }
        printf("\n");
    }

    printf("bench_dsa_select: done (lower ns is better; speedup = old/new)\n");
    free(isc); free(dst);
    return 0;
}