/* Synthetic finite-logit fixtures for pinned Colibri source.
 * Compile with -I /path/to/colibri/c. Header implementations are unchanged.
 * qsort is instrumented; rand supplies explicit endpoints and a midpoint.
 * This tests sampler behavior, not model inference or decode throughput. */
#include <assert.h>
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static size_t sorted_count;
static int sort_calls, draw_value, draw_calls;
static void observed_sort(void *base, size_t n, size_t width,
                          int (*compare)(const void *, const void *)) {
    sorted_count = n;
    sort_calls++;
    qsort(base, n, width, compare);
}
static int scripted_rand(void) { draw_calls++; return draw_value; }
#define qsort observed_sort
#define rand scripted_rand
#include "qwen36_sample.h"
#undef rand
#undef qsort

int main(void) {
    enum { V = 4096 };
    float logits[V];
    const char *names[] = {"flat", "spike", "plateau", "tail"};
    const float tops[] = {0.3f, 0.9f, 1.0f};
    const int draws[] = {0, RAND_MAX / 2, RAND_MAX};
    int comparisons = 0;
    for (int shape = 0; shape < 4; shape++) {
        for (int i = 0; i < V; i++) logits[i] = shape == 0 ? 0.f : -100.f;
        if (shape == 1) logits[1337] = 20.f;
        if (shape == 2) for (int i = 0; i < 50; i++) logits[(i * 79 + 13) % V] = 20.f;
        if (shape == 3) for (int i = 0; i < V; i++) logits[i] = -(float)(i % 83) / 3.f;
        for (int t = 0; t < 3; t++) {
            size_t candidates = 0;
            for (int d = 0; d < 3; d++) {
                draw_value = draws[d]; draw_calls = 0; sort_calls = 0;
                int got = serve_sample(logits, V, 0.7f, tops[t]);
                assert(draw_calls == 1 && sort_calls == 1);
                candidates = sorted_count;
                int want = serve_sample_sorted(logits, V, 0.7f, tops[t]);
                assert(sorted_count == V && draw_calls == 2 && sort_calls == 2);
                assert(got == want);
                assert(got >= 0 && got < V);
                if (shape == 1) assert(got == 1337);
                comparisons++;
            }
            if (t == 1 || (shape == 1 && t == 2))
                printf("%s: V=%d T=0.7 top_p=%.1f sorted_candidates=%zu/%d\n",
                       names[shape], V, tops[t], candidates, V);
        }
    }
    for (int i = 0; i < V; i++) logits[i] = 0.f;
    draw_value = RAND_MAX;
    int tie_last = serve_sample(logits, V, 0.7f, 0.9f);
    assert(tie_last == 3686);
    draw_value = 0;
    int tie_first = serve_sample(logits, V, 0.7f, 0.9f);
    assert(tie_first == 0);
    printf("flat_ties: draw=0 picks %d; draw=RAND_MAX picks %d; nucleus has 3687 token IDs\n", tie_first, tie_last);

    float special[33];
    for (int i = 0; i < 33; i++) special[i] = -1.f;
    special[17] = 9.f; special[32] = 9.f;
    sort_calls = 0; draw_calls = 0;
    assert(serve_sample(special, 33, 0.f, 0.9f) == 17);
    assert(sort_calls == 0 && draw_calls == 0);
    special[1] = NAN;
    assert(serve_sample(special, 33, 0.f, 0.9f) == 17);
    special[0] = NAN;
    assert(serve_sample(special, 33, 0.f, 0.9f) == 0);
    float zeros[] = {-0.f, 0.f, -0.f};
    assert(serve_sample(zeros, 3, 0.f, 0.9f) == 0);
    printf("greedy: no sort or draw; first maximum wins; later NaN ignored; initial NaN retains index 0; signed zeros tie\n");
    assert(comparisons == 36);
    printf("PASS: %d fast/full-sort comparisons, 2 tie endpoints, 4 greedy cases; synthetic logits only\n", comparisons);
    return 0;
}