#version 450

#extension GL_EXT_control_flow_attributes : enable
#extension GL_EXT_shader_16bit_storage : require

#include "types.glsl"

layout(constant_id = 0) const int BLOCK_SIZE = 1024;
layout(constant_id = 1) const int QSA = 0; // 1: fuse the qwen4 QSA indexer gather + f16 mask

layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;

layout (binding = 0) readonly buffer A  {float     data_a[];};   // input values, or QSA block scores [n_tps, n_blocks, n_stream]
layout (binding = 1) writeonly buffer D {int       data_d[];};   // [k, ...]
layout (binding = 2) readonly buffer CB {int       cell_blk[];}; // QSA: cell->block map [n_kv, n_stream]
layout (binding = 3) readonly buffer M  {float16_t mask[];};     // QSA: raw f16 kq_mask [n_kv, n_tps, n_stream]
layout (binding = 4) buffer S           {float     scratch[];};  // QSA: [nrows, n_kv] gathered inputs

layout (push_constant) uniform parameter {
    uint ncols;
    uint k;
    uint nrows;
    uint n_tps;    // QSA only
    uint n_blocks; // QSA only
    uint n_stream; // QSA only
} p;

#define RADIX_BITS 8
#define RADIX_SIZE (1 << RADIX_BITS)

shared uint histo[RADIX_SIZE];
shared uint sh_bucket;
shared uint sh_above;
shared uint out_count;

// order-preserving float -> uint mapping
uint f2ui(float x) {
    uint y = floatBitsToUint(x);
    if ((y & 0x80000000u) != 0u) {
        y ^= 0xFFFFFFFFu;
    } else {
        y |= 0x80000000u;
    }
    return y;
}

// QSA element i of row (t,s): score[cell_blk[i,s], t, s] + mask[i,t,s]
float gather(uint row, uint i) {
    const uint t = row % p.n_tps;
    const uint s = row / p.n_tps;
    const uint block = uint(cell_blk[s * p.ncols + i]);
    const float a = data_a[(s * p.n_blocks + block) * p.n_tps + t];
    const float m = float(mask[(s * p.n_tps + t) * p.ncols + i]);
    return a + m;
}

float load(uint row, uint i, bool first) {
    if (QSA == 0) {
        return data_a[row * p.ncols + i];
    }
    // materialize the scattered gather on the first pass and reuse it after; each
    // invocation only touches its own scratch entries, so no barrier is needed
    const uint off = row * p.ncols + i;
    if (first) {
        const float v = gather(row, i);
        scratch[off] = v;
        return v;
    }
    return scratch[off];
}

// one workgroup per row: radix-select the K-th largest, then compact it plus enough ties
void topk(const uint row) {
    const uint tid = gl_LocalInvocationID.x;
    const uint ncols = p.ncols;
    const uint row_out = row * p.k;

    uint prefix  = 0;     // fixed high bits of the threshold key
    uint desired = p.k;   // count still needed from the candidate range

    [[unroll]] for (int shift = 32 - RADIX_BITS; shift >= 0; shift -= RADIX_BITS) {
        for (uint i = tid; i < RADIX_SIZE; i += BLOCK_SIZE) {
            histo[i] = 0;
        }
        barrier();

        const bool first     = (shift == 32 - RADIX_BITS);
        const uint hi_mask   = (shift + RADIX_BITS >= 32) ? 0u : (0xFFFFFFFFu << uint(shift + RADIX_BITS));
        const uint prefix_hi = prefix & hi_mask;
        for (uint i = tid; i < ncols; i += BLOCK_SIZE) {
            const uint key = f2ui(load(row, i, first));
            if ((key & hi_mask) == prefix_hi) {
                atomicAdd(histo[(key >> uint(shift)) & (RADIX_SIZE - 1)], 1u);
            }
        }
        barrier();

        // top-down scan for the bucket holding the K-th value
        if (tid == 0) {
            uint acc = 0;
            uint b   = 0;
            for (int bb = RADIX_SIZE - 1; bb >= 0; --bb) {
                const uint c = histo[bb];
                if (acc + c >= desired) { b = uint(bb); break; }
                acc += c;
            }
            sh_bucket = b;
            sh_above  = acc;
        }
        barrier();

        prefix  |= sh_bucket << uint(shift);
        desired -= sh_above;
        barrier();
    }

    if (tid == 0) {
        out_count = 0;
    }
    barrier();

    // emit everything above the threshold, then fill the rest from ties
    const uint threshold = prefix;
    for (uint i = tid; i < ncols; i += BLOCK_SIZE) {
        if (f2ui(load(row, i, false)) > threshold) {
            data_d[row_out + atomicAdd(out_count, 1u)] = int(i);
        }
    }
    barrier();
    for (uint i = tid; i < ncols; i += BLOCK_SIZE) {
        if (f2ui(load(row, i, false)) == threshold) {
            const uint pos = atomicAdd(out_count, 1u);
            if (pos < p.k) {
                data_d[row_out + pos] = int(i);
            }
        }
    }
}

void main() {
    for (uint row = gl_WorkGroupID.y; row < p.nrows; row += gl_NumWorkGroups.y) {
        topk(row);
    }
}