#include "gem/sampler.h"

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

int main(void)
{
    char error[256];
    struct gem_sampler_config greedy_config = {
        .temperature = 0.0f, .top_k = 0, .top_p = 1.0f,
        .repeat_penalty = 2.0f, .repeat_last_n = 8, .seed = 1,
    };
    struct gem_sampler *sampler = NULL;
    if (gem_sampler_open(&sampler, 5, &greedy_config, error, sizeof(error)) != 0) {
        fprintf(stderr, "%s\n", error);
        return 1;
    }
    const float logits[] = {0.0f, 4.0f, 3.0f, -1.0f, 1.0f};
    const uint32_t history[] = {1, 1};
    uint32_t token;
    if (gem_sampler_sample(sampler, logits, history, 2, &token, error, sizeof(error)) != 0 ||
        token != 2) {
        fprintf(stderr, "repeat-penalty greedy mismatch: token %u (%s)\n", token, error);
        return 1;
    }
    gem_sampler_close(sampler);

    struct gem_sampler_config random_config = {
        .temperature = 0.8f, .top_k = 2, .top_p = 0.95f,
        .repeat_penalty = 1.0f, .repeat_last_n = 0, .seed = 42,
    };
    if (gem_sampler_open(&sampler, 5, &random_config, error, sizeof(error)) != 0)
        return 1;
    for (size_t i = 0; i < 32; ++i) {
        if (gem_sampler_sample(sampler, logits, NULL, 0, &token, error, sizeof(error)) != 0 ||
            (token != 1 && token != 2)) {
            fprintf(stderr, "top-k sample escaped candidate set\n");
            return 1;
        }
    }
    gem_sampler_close(sampler);

    const float parity_logits[] = {4.0f, 3.0f, 2.0f, -INFINITY, -INFINITY,
                                   -INFINITY, -INFINITY, -INFINITY};
    struct gem_sampler_config heap_config = random_config;
    heap_config.top_k = 3;
    heap_config.seed = 1234;
    struct gem_sampler_config sort_config = heap_config;
    sort_config.top_k = 0;
    struct gem_sampler *heap_sampler = NULL;
    struct gem_sampler *sort_sampler = NULL;
    if (gem_sampler_open(&heap_sampler, 8, &heap_config, error, sizeof(error)) != 0 ||
        gem_sampler_open(&sort_sampler, 8, &sort_config, error, sizeof(error)) != 0)
        return 1;
    for (size_t i = 0; i < 128; ++i) {
        uint32_t heap_token, sort_token;
        if (gem_sampler_sample(heap_sampler, parity_logits, NULL, 0, &heap_token,
                               error, sizeof(error)) != 0 ||
            gem_sampler_sample(sort_sampler, parity_logits, NULL, 0, &sort_token,
                               error, sizeof(error)) != 0 ||
            heap_token != sort_token) {
            fprintf(stderr, "heap/full-sort parity mismatch at sample %zu\n", i);
            return 1;
        }
    }
    gem_sampler_close(sort_sampler);
    gem_sampler_close(heap_sampler);

    const float tied_logits[] = {0.0f, 4.0f, 4.0f, 4.0f};
    heap_config.top_k = 1;
    if (gem_sampler_open(&sampler, 4, &heap_config, error, sizeof(error)) != 0 ||
        gem_sampler_sample(sampler, tied_logits, NULL, 0, &token,
                           error, sizeof(error)) != 0 || token != 1) {
        fprintf(stderr, "deterministic top-k tie mismatch: token %u\n", token);
        return 1;
    }
    gem_sampler_close(sampler);

    struct gem_sampler_config distribution_config = {
        .temperature = 1.0f, .top_k = 3, .top_p = 1.0f,
        .repeat_penalty = 1.0f, .repeat_last_n = 0, .seed = 7,
    };
    const float distribution_logits[] = {0.0f, logf(2.0f), -INFINITY, logf(3.0f)};
    double probabilities[4];
    if (gem_sampler_open(&sampler, 4, &distribution_config, error, sizeof(error)) != 0 ||
        gem_sampler_distribution(sampler, distribution_logits, NULL, 0,
                                 probabilities, error, sizeof(error)) != 0 ||
        fabs(probabilities[0] - 1.0 / 6.0) > 1e-7 ||
        fabs(probabilities[1] - 2.0 / 6.0) > 1e-7 || probabilities[2] != 0.0 ||
        fabs(probabilities[3] - 3.0 / 6.0) > 1e-7) {
        fprintf(stderr, "filtered sampler distribution mismatch: %s\n", error);
        return 1;
    }
    uint32_t drawn;
    bool accepted;
    if (gem_sampler_draw_distribution(sampler, probabilities, &drawn,
                                      error, sizeof(error)) != 0 || drawn >= 4 ||
        gem_sampler_accept(sampler, 0.0, &accepted, error, sizeof(error)) != 0 || accepted ||
        gem_sampler_accept(sampler, 1.0, &accepted, error, sizeof(error)) != 0 || !accepted) {
        fprintf(stderr, "probability draw/accept mismatch: %s\n", error);
        return 1;
    }
    const double target_probabilities[] = {0.25, 0.75};
    const double draft_probabilities[] = {0.5, 0.5};
    double residual[2];
    if (fabs(gem_sampler_distribution_overlap(target_probabilities,
                                               draft_probabilities, 2) - 0.75) > 1e-12 ||
        gem_sampler_positive_difference(target_probabilities, draft_probabilities, 2,
                                        residual, error, sizeof(error)) != 0 ||
        residual[0] != 0.0 || residual[1] != 1.0) {
        fprintf(stderr, "speculative residual distribution mismatch: %s\n", error);
        return 1;
    }
    gem_sampler_close(sampler);
    return 0;
}