#include "gem/cpu.h"

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

static uint8_t *zero_q4_matrix(size_t input_size, size_t output_size)
{
    size_t bytes = input_size / 32 * 18 * output_size;
    uint8_t *matrix = calloc(bytes, 1);
    if (!matrix)
        return NULL;
    for (size_t offset = 0; offset < bytes; offset += 18) {
        matrix[offset] = 0x00;
        matrix[offset + 1] = 0x3c; /* F16 1.0 */
        memset(matrix + offset + 2, 0x88, 16);
    }
    return matrix;
}

int main(void)
{
    uint8_t *embedding_matrix = zero_q4_matrix(32, 32);
    uint8_t *kv_matrix = zero_q4_matrix(32, 16);
    if (!embedding_matrix || !kv_matrix)
        return 2;
    float norm32[32];
    float norm16[16];
    float input[32];
    float output[32];
    for (size_t i = 0; i < 32; ++i) {
        norm32[i] = 1.0f;
        input[i] = (float)i - 16.0f;
    }
    for (size_t i = 0; i < 16; ++i)
        norm16[i] = 1.0f;
    struct gem_cpu_layer layer = {
        .embedding_size = 32,
        .query_heads = 2,
        .kv_heads = 1,
        .head_dimension = 16,
        .rotary_dimension = 16,
        .feed_forward_size = 32,
        .rms_epsilon = 1e-6f,
        .rope_frequency_base = 10000.0f,
        .rope_frequency_scale = 1.0f,
        .attention_scale = 1.0f,
        .attention_query = embedding_matrix,
        .attention_key = kv_matrix,
        .attention_value = NULL, /* exercise K-as-V */
        .attention_output = embedding_matrix,
        .ffn_gate = embedding_matrix,
        .ffn_up = embedding_matrix,
        .ffn_down = embedding_matrix,
        .attention_norm = norm32,
        .query_norm = norm16,
        .key_norm = norm16,
        .post_attention_norm = norm32,
        .ffn_norm = norm32,
        .post_ffn_norm = norm32,
        .layer_output_scale = 0.5f,
    };
    float key_cache[4 * 16] = {0};
    float value_cache[4 * 16] = {0};
    char error[256];
    if (gem_cpu_layer_decode(&layer, input, 0, key_cache, value_cache, 4, 0, 0,
                             output, error, sizeof(error)) != 0) {
        fprintf(stderr, "%s\n", error);
        return 1;
    }
    for (size_t i = 0; i < 32; ++i)
        if (fabsf(output[i] - input[i] * 0.5f) > 1e-6f) {
            fprintf(stderr, "layer output mismatch at %zu: %.8f\n", i, output[i]);
            return 1;
        }
    free(kv_matrix);
    free(embedding_matrix);
    return 0;
}