#ifndef DIM
#define DIM 128
#endif

#ifndef TOPK
#define TOPK 32
#endif

#ifndef QUERY_COUNT
#define QUERY_COUNT 1
#endif

/*
 * Numero di work-item che collaborano per ogni query.
 */
#ifndef LANES_PER_QUERY
#define LANES_PER_QUERY 8
#endif

/*
 * Ogni work-group contiene 8 work-item.
 *
 * Ogni work-item calcola un TOPK locale.
 * Poi il work-item 0 fonde gli 8 TOPK locali.
 */

__kernel void knn(
    __global const float *db,
    __global const float *queries,
    __global uint *indices,
    __global float *values,
    uint n)
{
    uint lane = get_local_id(0);
    uint q = get_group_id(0);

    if (q >= QUERY_COUNT)
        return;

    /*
     * TOPK locale per ogni work-item.
     */
    float best_score[TOPK];
    uint best_index[TOPK];

    for (uint k = 0; k < TOPK; ++k) {
        best_score[k] = -3.402823466e+38f;
        best_index[k] = 0xffffffff;
    }

    __global const float *query =
        &queries[q * DIM];

    /*
     * Ogni lane analizza una parte del database.
     *
     * Esempio con 8 lane:
     *
     * lane 0 -> 0, 8, 16, 24, ...
     * lane 1 -> 1, 9, 17, 25, ...
     * ...
     * lane 7 -> 7, 15, 23, 31, ...
     */
    for (uint i = lane; i < n; i += LANES_PER_QUERY) {

        __global const float *v =
            &db[i * DIM];

        float s = 0.0f;

        for (uint j = 0; j < DIM; ++j)
            s += v[j] * query[j];

        if (s <= best_score[TOPK - 1])
            continue;

        int pos = TOPK - 1;

        while (pos > 0 &&
               s > best_score[pos - 1]) {
            --pos;
        }

        for (int k = TOPK - 1; k > pos; --k) {
            best_score[k] =
                best_score[k - 1];

            best_index[k] =
                best_index[k - 1];
        }

        best_score[pos] = s;
        best_index[pos] = i;
    }

    /*
     * TOPK locali nella local memory.
     *
     * 8 lane × 32 elementi.
     */
    __local float local_scores[LANES_PER_QUERY][TOPK];
    __local uint local_indices[LANES_PER_QUERY][TOPK];

    for (uint k = 0; k < TOPK; ++k) {
        local_scores[lane][k] = best_score[k];
        local_indices[lane][k] = best_index[k];
    }

    /*
     * Tutti i lane devono aver scritto
     * la local memory prima della riduzione.
     */
    barrier(CLK_LOCAL_MEM_FENCE);

    /*
     * Solo lane 0 esegue il merge finale.
     */
    if (lane == 0) {

        float final_score[TOPK];
        uint final_index[TOPK];

        for (uint k = 0; k < TOPK; ++k) {
            final_score[k] =
                -3.402823466e+38f;

            final_index[k] =
                0xffffffff;
        }

        /*
         * Fonde gli 8 TOPK locali.
         *
         * Al massimo 8 × 32 = 256 candidati.
         */
        for (uint l = 0;
             l < LANES_PER_QUERY;
             ++l) {

            for (uint k = 0;
                 k < TOPK;
                 ++k) {

                float s =
                    local_scores[l][k];

                uint idx =
                    local_indices[l][k];

                if (idx == 0xffffffff)
                    continue;

                if (s <= final_score[TOPK - 1])
                    continue;

                int pos = TOPK - 1;

                while (pos > 0 &&
                       s > final_score[pos - 1]) {
                    --pos;
                }

                for (int p = TOPK - 1;
                     p > pos;
                     --p) {

                    final_score[p] =
                        final_score[p - 1];

                    final_index[p] =
                        final_index[p - 1];
                }

                final_score[pos] = s;
                final_index[pos] = idx;
            }
        }

        /*
         * Risultato finale della query.
         */
        for (uint k = 0; k < TOPK; ++k) {
            indices[q * TOPK + k] =
                final_index[k];

            values[q * TOPK + k] =
                final_score[k];
        }
    }
}