#include "gem/ops.h"

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

static int close_enough(float actual, float expected, float tolerance)
{
    if (fabsf(actual - expected) <= tolerance)
        return 1;
    fprintf(stderr, "expected %.8f, got %.8f\n", expected, actual);
    return 0;
}

static int test_online_softmax_partitions(void)
{
    enum { TOKENS = 13, DIMENSIONS = 17, MAX_PARTITIONS = 16 };
    float scores[TOKENS];
    float values[TOKENS * DIMENSIONS];
    float reference_numerator[DIMENSIONS];
    float reference[DIMENSIONS];
    float partial_numerator[MAX_PARTITIONS][DIMENSIONS];
    float partial_maximum[MAX_PARTITIONS];
    float partial_denominator[MAX_PARTITIONS];
    float merged_numerator[DIMENSIONS];
    float scratch_numerator[DIMENSIONS];
    float actual[DIMENSIONS];
    for (size_t token = 0; token < TOKENS; ++token) {
        scores[token] = ((int)(token * 19 % 23) - 11) * 0.75f;
        for (size_t dimension = 0; dimension < DIMENSIONS; ++dimension)
            values[token * DIMENSIONS + dimension] =
                ((int)((token * 31 + dimension * 7) % 41) - 20) * 0.125f;
    }
    float reference_maximum;
    float reference_denominator;
    if (gem_attention_online_partial_f32(
            scores, values, TOKENS, DIMENSIONS, &reference_maximum,
            &reference_denominator, reference_numerator) != 0 ||
        gem_attention_online_finalize_f32(
            reference_denominator, reference_numerator, DIMENSIONS, reference) != 0)
        return 0;

    static const size_t partition_counts[] = {1, 2, 3, 4, 7, 16};
    for (size_t test = 0;
         test < sizeof(partition_counts) / sizeof(partition_counts[0]); ++test) {
        size_t partitions = partition_counts[test];
        for (size_t partition = 0; partition < partitions; ++partition) {
            size_t begin = TOKENS * partition / partitions;
            size_t end = TOKENS * (partition + 1) / partitions;
            if (gem_attention_online_partial_f32(
                    scores + begin, values + begin * DIMENSIONS, end - begin, DIMENSIONS,
                    &partial_maximum[partition], &partial_denominator[partition],
                    partial_numerator[partition]) != 0)
                return 0;
        }
        float merged_maximum = -INFINITY;
        float merged_denominator = 0.0f;
        memset(merged_numerator, 0, sizeof(merged_numerator));
        for (size_t partition = 0; partition < partitions; ++partition) {
            float next_maximum;
            float next_denominator;
            if (gem_attention_online_merge_f32(
                    merged_maximum, merged_denominator, merged_numerator,
                    partial_maximum[partition], partial_denominator[partition],
                    partial_numerator[partition], DIMENSIONS,
                    &next_maximum, &next_denominator, scratch_numerator) != 0)
                return 0;
            merged_maximum = next_maximum;
            merged_denominator = next_denominator;
            memcpy(merged_numerator, scratch_numerator, sizeof(merged_numerator));
        }
        if (gem_attention_online_finalize_f32(
                merged_denominator, merged_numerator, DIMENSIONS, actual) != 0)
            return 0;
        for (size_t dimension = 0; dimension < DIMENSIONS; ++dimension)
            if (!close_enough(actual[dimension], reference[dimension], 2e-6f))
                return 0;
    }

    float invalid_maximum;
    float invalid_denominator;
    float invalid_numerator[DIMENSIONS];
    float invalid_scores[] = {NAN};
    size_t maximum_bytes;
    size_t denominator_bytes;
    size_t numerator_bytes;
    size_t total_bytes;
    return gem_attention_online_partial_f32(
               NULL, NULL, 0, DIMENSIONS, &invalid_maximum,
               &invalid_denominator, invalid_numerator) == 0 &&
           isinf(invalid_maximum) && invalid_maximum < 0.0f &&
           invalid_denominator == 0.0f &&
           gem_attention_online_finalize_f32(
               invalid_denominator, invalid_numerator, DIMENSIONS, actual) != 0 &&
           gem_attention_online_partial_f32(
               invalid_scores, values, 1, DIMENSIONS, &invalid_maximum,
               &invalid_denominator, invalid_numerator) != 0 &&
           gem_attention_online_scratch_bytes(
               32, 512, 32, &maximum_bytes, &denominator_bytes,
               &numerator_bytes, &total_bytes) == 0 &&
           maximum_bytes == 4096 && denominator_bytes == 4096 &&
           numerator_bytes == 2097152 && total_bytes == 2105344 &&
           gem_attention_online_scratch_bytes(
               SIZE_MAX, 512, 32, &maximum_bytes, &denominator_bytes,
               &numerator_bytes, &total_bytes) != 0;
}

int main(void)
{
    if (!test_online_softmax_partitions())
        return 1;
    if (!close_enough(gem_f16_to_f32(UINT16_C(0x3c00)), 1.0f, 0.0f) ||
        !close_enough(gem_f16_to_f32(UINT16_C(0xc000)), -2.0f, 0.0f) ||
        !close_enough(gem_f16_to_f32(UINT16_C(0x0001)), 0x1p-24f, 0.0f))
        return 1;

    uint8_t q4[18] = {0x00, 0x3c};
    memset(q4 + 2, 0xf0, 16);
    float values[32];
    if (gem_dequantize_q4_0(q4, 32, values) != 0)
        return 1;
    for (size_t i = 0; i < 16; ++i)
        if (!close_enough(values[i], -8.0f, 0.0f) ||
            !close_enough(values[i + 16], 7.0f, 0.0f))
            return 1;

    float input[32];
    for (size_t i = 0; i < 32; ++i)
        input[i] = 1.0f;
    float product;
    if (gem_matvec_q4_0(q4, 32, 1, input, &product) != 0 ||
        !close_enough(product, -16.0f, 0.0f))
        return 1;

    uint8_t q8[34] = {0x00, 0x38}; /* scale 0.5 */
    for (size_t i = 0; i < 32; ++i)
        q8[2 + i] = (uint8_t)(int8_t)((int)i - 16);
    if (gem_dequantize_q8_0(q8, 32, values) != 0)
        return 1;
    for (size_t i = 0; i < 32; ++i)
        if (!close_enough(values[i], ((float)i - 16.0f) * 0.5f, 0.0f))
            return 1;
    if (gem_matvec_q8_0(q8, 32, 1, input, &product) != 0 ||
        !close_enough(product, -8.0f, 0.0f))
        return 1;

    const float norm_input[] = {3.0f, 4.0f};
    const float weights[] = {1.0f, 2.0f};
    float norm_output[2];
    gem_rms_norm_f32(norm_input, weights, 0.0f, 2, norm_output);
    float inverse_rms = 1.0f / sqrtf(12.5f);
    if (!close_enough(norm_output[0], 3.0f * inverse_rms, 1e-6f) ||
        !close_enough(norm_output[1], 8.0f * inverse_rms, 1e-6f))
        return 1;

    float softcap[] = {-60.0f, 0.0f, 60.0f};
    gem_softcap_f32(softcap, 3, 30.0f);
    if (!close_enough(softcap[0], -tanhf(2.0f) * 30.0f, 1e-6f) ||
        !close_enough(softcap[1], 0.0f, 0.0f) ||
        !close_enough(softcap[2], tanhf(2.0f) * 30.0f, 1e-6f))
        return 1;

    const float rope_input[] = {1.0f, 2.0f, 3.0f, 4.0f};
    float rope_output[4];
    if (gem_rope_neox_f32(rope_input, 4, 4, 1, 10000.0f, 1.0f, NULL, rope_output) != 0 ||
        !close_enough(rope_output[0], cosf(1.0f) - 3.0f * sinf(1.0f), 1e-6f) ||
        !close_enough(rope_output[2], sinf(1.0f) + 3.0f * cosf(1.0f), 1e-6f) ||
        !close_enough(rope_output[1], 2.0f * cosf(0.01f) - 4.0f * sinf(0.01f), 1e-6f) ||
        !close_enough(rope_output[3], 2.0f * sinf(0.01f) + 4.0f * cosf(0.01f), 1e-6f))
        return 1;

    const float query[] = {1.0f, 0.0f, 0.0f, 1.0f};
    const float keys[] = {1.0f, 0.0f, 0.0f, 1.0f};
    const float values_cache[] = {10.0f, 20.0f, 30.0f, 40.0f};
    float scores[2];
    float attention[4];
    float p = expf(1.0f) / (expf(1.0f) + 1.0f);
    if (gem_attention_decode_f32(query, keys, values_cache, 2, 1, 2, 2, 0, 1.0f,
                                 scores, attention) != 0 ||
        !close_enough(attention[0], p * 10.0f + (1.0f - p) * 30.0f, 1e-5f) ||
        !close_enough(attention[1], p * 20.0f + (1.0f - p) * 40.0f, 1e-5f) ||
        !close_enough(attention[2], (1.0f - p) * 10.0f + p * 30.0f, 1e-5f) ||
        !close_enough(attention[3], (1.0f - p) * 20.0f + p * 40.0f, 1e-5f))
        return 1;
    if (gem_attention_decode_f32(query, keys, values_cache, 2, 1, 2, 2, 1, 1.0f,
                                 scores, attention) != 0)
        return 1;
    for (size_t i = 0; i < 4; i += 2)
        if (!close_enough(attention[i], 30.0f, 0.0f) ||
            !close_enough(attention[i + 1], 40.0f, 0.0f))
            return 1;

    const float ring_keys[] = {0.0f, 1.0f, 1.0f, 0.0f}; /* logical t2 in slot 0, t1 in slot 1 */
    const float ring_values[] = {30.0f, 40.0f, 10.0f, 20.0f};
    if (gem_attention_decode_ring_f32(query, ring_keys, ring_values, 2, 1, 2, 3, 2, 2,
                                      1.0f, scores, attention) != 0 ||
        !close_enough(attention[0], p * 10.0f + (1.0f - p) * 30.0f, 1e-5f) ||
        !close_enough(attention[3], p * 40.0f + (1.0f - p) * 20.0f, 1e-5f))
        return 1;
    return 0;
}