#include <string.h>

#if defined(__AVX2__)
#include <immintrin.h>
#endif

#include "bio_quant.h"

float bio_quant_decode_sample(const uint8_t *data, uint64_t len, uint64_t index,
                              uint32_t ggml_type) {
    const uint8_t b = data == NULL || len == 0u ? 0u : data[index % len];
    const uint8_t nibble = (index & 1u) != 0u ? (uint8_t)((b >> 4) & 0x0fu)
                                             : (uint8_t)(b & 0x0fu);
    const int centered = (int)nibble - 8;
    switch (ggml_type) {
    case BIO_GGUF_TYPE_Q4_K:
    case BIO_GGUF_TYPE_Q4_0:
    case BIO_GGUF_TYPE_Q4_1:
        return (float)centered / 8.0f;
    case BIO_GGUF_TYPE_Q6_K:
        return (float)((int)(b & 0x3fu) - 32) / 32.0f;
    default:
        return (float)b / 255.0f;
    }
}

bio_v1_status bio_quant_project_scalar(const bio_gguf_tensor_span *tensor,
                                       const uint32_t *tokens,
                                       uint32_t token_count,
                                       float *logits,
                                       uint32_t logit_count) {
    if (tensor == NULL || tensor->data == NULL || tokens == NULL || logits == NULL ||
        token_count == 0u || logit_count == 0u) {
        return BIO_ERR_ARGS;
    }
    for (uint32_t j = 0; j < logit_count; ++j) {
        float acc = 0.0f;
        for (uint32_t i = 0; i < token_count; ++i) {
            const uint64_t index = ((uint64_t)tokens[i] * 1315423911ull) + j * 2654435761ull + i;
            acc += bio_quant_decode_sample(tensor->data, tensor->byte_size, index,
                                           tensor->ggml_type);
        }
        logits[j] = acc;
    }
    return BIO_OK;
}

bio_v1_status bio_quant_project_avx2(const bio_gguf_tensor_span *tensor,
                                     const uint32_t *tokens,
                                     uint32_t token_count,
                                     float *logits,
                                     uint32_t logit_count) {
#if defined(__AVX2__)
    uint32_t j = 0;
    if (tensor == NULL || tensor->data == NULL || tokens == NULL || logits == NULL ||
        token_count == 0u || logit_count == 0u) {
        return BIO_ERR_ARGS;
    }
    for (; j + 8u <= logit_count; j += 8u) {
        __m256 acc = _mm256_setzero_ps();
        for (uint32_t i = 0; i < token_count; ++i) {
            float tmp[8];
            for (uint32_t lane = 0; lane < 8u; ++lane) {
                const uint64_t index =
                    ((uint64_t)tokens[i] * 1315423911ull) +
                    (uint64_t)(j + lane) * 2654435761ull + i;
                tmp[lane] = bio_quant_decode_sample(tensor->data, tensor->byte_size, index,
                                                    tensor->ggml_type);
            }
            acc = _mm256_add_ps(acc, _mm256_loadu_ps(tmp));
        }
        _mm256_storeu_ps(logits + j, acc);
    }
    if (j < logit_count) {
        return bio_quant_project_scalar(tensor, tokens, token_count, logits + j,
                                        logit_count - j);
    }
    return BIO_OK;
#else
    return bio_quant_project_scalar(tensor, tokens, token_count, logits, logit_count);
#endif
}
