#include <string.h>

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

#include "bio_kernel.h"

bio_kernel_dispatch g_bio_kernels = {
    "scalar_ref",
    bio_pulse_int4_scalar,
};

void bio_runtime_bind_kernels(void) {
    const uint32_t features = bio_v1_cpu_features();
#if defined(__AVX2__)
    if ((features & BIO_CPU_AVX2) != 0u) {
        g_bio_kernels.pulse_variant = "avx2";
        g_bio_kernels.pulse_int4 = bio_pulse_int4_avx2;
        return;
    }
#endif
#if defined(__aarch64__) || defined(_M_ARM64)
    if ((features & BIO_CPU_NEON) != 0u) {
        g_bio_kernels.pulse_variant = "neon_pending_scalar_ref";
        g_bio_kernels.pulse_int4 = bio_pulse_int4_scalar;
        return;
    }
#endif
    g_bio_kernels.pulse_variant = "scalar_ref";
    g_bio_kernels.pulse_int4 = bio_pulse_int4_scalar;
}

static int block_is_active(const bio_sparse_masks *masks, size_t block_index) {
    if (masks == NULL || masks->block_count == 0u) {
        return 1;
    }
    if (block_index >= masks->block_count) {
        return 0;
    }
    return masks->omega_active[block_index] != 0u && masks->zpm_skip[block_index] == 0u;
}

bio_v1_status bio_pulse_int4_scalar(uint8_t *packed_potentials,
                                    const uint8_t *packed_inputs,
                                    const bio_sparse_masks *masks,
                                    size_t byte_count,
                                    uint8_t threshold,
                                    uint64_t *executed_updates_out,
                                    uint64_t *baseline_updates_out) {
    const size_t block_bytes =
        (masks != NULL && masks->block_bytes != 0u) ? masks->block_bytes
                                                    : BIO_ROUTE_DEFAULT_BLOCK_BYTES;
    uint64_t executed = 0;
    uint64_t baseline = (uint64_t)byte_count * 2u;
    if (packed_potentials == NULL || packed_inputs == NULL || threshold == 0u ||
        threshold > 15u) {
        return BIO_ERR_ARGS;
    }
    for (size_t off = 0; off < byte_count; off += block_bytes) {
        const size_t block = off / block_bytes;
        const size_t end = (off + block_bytes <= byte_count) ? off + block_bytes : byte_count;
        if (!block_is_active(masks, block)) {
            continue;
        }
        for (size_t i = off; i < end; ++i) {
            const uint8_t p = packed_potentials[i];
            const uint8_t in = packed_inputs[i];
            uint8_t lo = (uint8_t)((p & 0x0fu) + (in & 0x0fu));
            uint8_t hi = (uint8_t)(((p >> 4) & 0x0fu) + ((in >> 4) & 0x0fu));
            if (lo > 15u) {
                lo = 15u;
            }
            if (hi > 15u) {
                hi = 15u;
            }
            if (lo >= threshold) {
                lo = 0u;
            }
            if (hi >= threshold) {
                hi = 0u;
            }
            packed_potentials[i] = (uint8_t)(lo | (uint8_t)(hi << 4));
            executed += 2u;
        }
    }
    if (executed_updates_out != NULL) {
        *executed_updates_out = executed;
    }
    if (baseline_updates_out != NULL) {
        *baseline_updates_out = baseline;
    }
    return BIO_OK;
}

bio_v1_status bio_pulse_int4_avx2(uint8_t *packed_potentials,
                                  const uint8_t *packed_inputs,
                                  const bio_sparse_masks *masks,
                                  size_t byte_count,
                                  uint8_t threshold,
                                  uint64_t *executed_updates_out,
                                  uint64_t *baseline_updates_out) {
#if defined(__AVX2__)
    const size_t block_bytes =
        (masks != NULL && masks->block_bytes != 0u) ? masks->block_bytes
                                                    : BIO_ROUTE_DEFAULT_BLOCK_BYTES;
    const __m256i nibble_mask = _mm256_set1_epi8(0x0f);
    const __m256i v_thresh = _mm256_set1_epi8((char)(threshold & 0x0f));
    uint64_t executed = 0;
    uint64_t baseline = (uint64_t)byte_count * 2u;
    if (packed_potentials == NULL || packed_inputs == NULL || threshold == 0u ||
        threshold > 15u) {
        return BIO_ERR_ARGS;
    }
    for (size_t off = 0; off < byte_count; off += block_bytes) {
        const size_t block = off / block_bytes;
        const size_t end = (off + block_bytes <= byte_count) ? off + block_bytes : byte_count;
        size_t i = off;
        if (!block_is_active(masks, block)) {
            continue;
        }
        for (; i + 32u <= end; i += 32u) {
            const __m256i v_pot =
                _mm256_loadu_si256((const __m256i *)&packed_potentials[i]);
            const __m256i v_in =
                _mm256_loadu_si256((const __m256i *)&packed_inputs[i]);
            __m256i lo_p = _mm256_and_si256(v_pot, nibble_mask);
            __m256i hi_p = _mm256_and_si256(_mm256_srli_epi16(v_pot, 4), nibble_mask);
            const __m256i lo_i = _mm256_and_si256(v_in, nibble_mask);
            const __m256i hi_i = _mm256_and_si256(_mm256_srli_epi16(v_in, 4), nibble_mask);
            lo_p = _mm256_add_epi8(lo_p, lo_i);
            hi_p = _mm256_add_epi8(hi_p, hi_i);
            lo_p = _mm256_min_epu8(lo_p, nibble_mask);
            hi_p = _mm256_min_epu8(hi_p, nibble_mask);
            {
                const __m256i m_lo =
                    _mm256_cmpeq_epi8(_mm256_max_epu8(lo_p, v_thresh), lo_p);
                const __m256i m_hi =
                    _mm256_cmpeq_epi8(_mm256_max_epu8(hi_p, v_thresh), hi_p);
                lo_p = _mm256_andnot_si256(m_lo, lo_p);
                hi_p = _mm256_andnot_si256(m_hi, hi_p);
            }
            _mm256_storeu_si256(
                (__m256i *)&packed_potentials[i],
                _mm256_or_si256(_mm256_and_si256(lo_p, nibble_mask),
                                _mm256_slli_epi16(_mm256_and_si256(hi_p, nibble_mask), 4)));
            executed += 64u;
        }
        for (; i < end; ++i) {
            const uint8_t p = packed_potentials[i];
            const uint8_t in = packed_inputs[i];
            uint8_t lo = (uint8_t)((p & 0x0fu) + (in & 0x0fu));
            uint8_t hi = (uint8_t)(((p >> 4) & 0x0fu) + ((in >> 4) & 0x0fu));
            if (lo > 15u) {
                lo = 15u;
            }
            if (hi > 15u) {
                hi = 15u;
            }
            if (lo >= threshold) {
                lo = 0u;
            }
            if (hi >= threshold) {
                hi = 0u;
            }
            packed_potentials[i] = (uint8_t)(lo | (uint8_t)(hi << 4));
            executed += 2u;
        }
    }
    if (executed_updates_out != NULL) {
        *executed_updates_out = executed;
    }
    if (baseline_updates_out != NULL) {
        *baseline_updates_out = baseline;
    }
    return BIO_OK;
#else
    return bio_pulse_int4_scalar(packed_potentials, packed_inputs, masks, byte_count,
                                 threshold, executed_updates_out, baseline_updates_out);
#endif
}
