#include <string.h>

#include "bio_kernel.h"

static uint64_t fnv1a_block(const uint8_t *p, size_t n) {
    uint64_t x = 1469598103934665603ull;
    for (size_t i = 0; i < n; ++i) {
        x ^= (uint64_t)p[i];
        x *= 1099511628211ull;
    }
    return x;
}

static uint64_t count_nonzero_bytes(const uint8_t *p, size_t n) {
    uint64_t c = 0;
    for (size_t i = 0; i < n; ++i) {
        c += p[i] != 0u ? 1u : 0u;
    }
    return c;
}

void bio_sparse_masks_reset(bio_sparse_masks *masks, size_t byte_count,
                            size_t block_bytes) {
    if (masks == NULL) {
        return;
    }
    memset(masks, 0, sizeof(*masks));
    masks->block_bytes = block_bytes == 0u ? BIO_ROUTE_DEFAULT_BLOCK_BYTES : block_bytes;
    masks->block_count =
        (byte_count + masks->block_bytes - 1u) / masks->block_bytes;
    if (masks->block_count > BIO_ROUTE_MAX_BLOCKS) {
        masks->block_count = BIO_ROUTE_MAX_BLOCKS;
    }
    for (size_t i = 0; i < masks->block_count; ++i) {
        masks->omega_active[i] = 1u;
    }
    masks->baseline_updates = (uint64_t)byte_count * 2u;
    masks->active_blocks = (uint64_t)masks->block_count;
}

void bio_sparse_zpm_build(bio_sparse_masks *masks, const uint8_t *packed_inputs,
                          size_t byte_count, uint64_t zero_threshold) {
    if (masks == NULL || packed_inputs == NULL || masks->block_bytes == 0u) {
        return;
    }
    masks->skipped_blocks = 0;
    for (size_t block = 0; block < masks->block_count; ++block) {
        const size_t off = block * masks->block_bytes;
        const size_t n = (off + masks->block_bytes <= byte_count)
                             ? masks->block_bytes
                             : (byte_count - off);
        const uint64_t nonzero = count_nonzero_bytes(packed_inputs + off, n);
        masks->zpm_skip[block] = nonzero <= zero_threshold ? 1u : 0u;
        masks->skipped_blocks += masks->zpm_skip[block] != 0u ? 1u : 0u;
    }
}

void bio_sparse_omega_build(bio_sparse_masks *masks, const uint8_t *packed_inputs,
                            size_t byte_count, uint32_t wake_mod,
                            size_t min_active_blocks) {
    size_t active = 0;
    const uint32_t safe_wake_mod = wake_mod == 0u ? 1u : wake_mod;
    if (masks == NULL || packed_inputs == NULL || masks->block_bytes == 0u) {
        return;
    }
    for (size_t block = 0; block < masks->block_count; ++block) {
        const size_t off = block * masks->block_bytes;
        const size_t n = (off + masks->block_bytes <= byte_count)
                             ? masks->block_bytes
                             : (byte_count - off);
        const uint64_t sig = fnv1a_block(packed_inputs + off, n);
        const int by_index = (block % (size_t)safe_wake_mod) == 0u;
        const int by_signature = (uint32_t)(sig % (uint64_t)safe_wake_mod) == 0u;
        masks->omega_active[block] =
            masks->zpm_skip[block] == 0u && (by_index || by_signature) ? 1u : 0u;
        active += masks->omega_active[block] != 0u ? 1u : 0u;
    }
    if (min_active_blocks > masks->block_count) {
        min_active_blocks = masks->block_count;
    }
    for (size_t block = 0; block < masks->block_count && active < min_active_blocks; ++block) {
        if (masks->zpm_skip[block] == 0u && masks->omega_active[block] == 0u) {
            masks->omega_active[block] = 1u;
            ++active;
        }
    }
    masks->active_blocks = (uint64_t)active;
}

void bio_sparse_prune(bio_sparse_masks *masks) {
    uint64_t active = 0;
    if (masks == NULL) {
        return;
    }
    for (size_t block = 0; block < masks->block_count; ++block) {
        if (masks->zpm_skip[block] != 0u) {
            masks->omega_active[block] = 0u;
        }
        active += masks->omega_active[block] != 0u ? 1u : 0u;
    }
    masks->active_blocks = active;
}
