#version 450
// Copyright (c) 2025 Hans-Kristian Arntzen
// SPDX-License-Identifier: MIT
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_ballot : require
#extension GL_KHR_shader_subgroup_arithmetic : require
#extension GL_KHR_shader_subgroup_shuffle_relative : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_KHR_shader_subgroup_clustered : require
#extension GL_KHR_shader_subgroup_vote : require
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
#extension GL_EXT_shader_8bit_storage : require

#include "constants.h"

layout(local_size_x = 64) in;

struct BlockMeta
{
    uint code_word;
    uint offset;
};

struct BitstreamPacket
{
    uint offset;
    uint num_words;
};

layout(set = 0, binding = 0) writeonly buffer BitstreamPayload
{
    uint data[];
} bitstream_data;

layout(set = 0, binding = 0) writeonly buffer BitstreamPayload16Bit
{
    uint16_t data[];
} bitstream_data_16b;

layout(set = 0, binding = 0) writeonly buffer BitstreamPayload8Bit
{
    uint8_t data[];
} bitstream_data_8b;

layout(set = 0, binding = 1) writeonly buffer BitstreamMeta
{
    BitstreamPacket packets[];
} bitstream_meta;

layout(set = 0, binding = 2) readonly buffer SSBOMeta
{
    BlockMeta meta[];
} block_meta;

layout(set = 0, binding = 3) buffer Payloads
{
    layout(offset = 4) uint bitstream_payload_counter;
    layout(offset = 8) uint8_t data[];
} payload_data;

struct QuantStats
{
    float16_t square_error;
    uint16_t payload_cost;
};

struct BlockStats
{
    uint num_planes;
    QuantStats errors[15];
};

layout(set = 0, binding = 4) readonly buffer SSBOBlockStats
{
    BlockStats stats[];
} block_stats;

layout(set = 0, binding = 5) readonly buffer RateControlQuant
{
    int data[];
} quant_data;

layout(push_constant) uniform Registers
{
    ivec2 resolution;
    ivec2 resolution_32x32_blocks;
    ivec2 resolution_8x8_blocks;
    uint quant_resolution_code;
    uint sequence_code;
    int block_offset_32x32;
    int block_stride_32x32;
    int block_offset_8x8;
    int block_stride_8x8;
} registers;

uint compute_required_8x8_size(uint control_word)
{
    int q_bits = int(bitfieldExtract(control_word, Q_PLANES_OFFSET, Q_PLANES_BITS));
    uint lsbs = control_word & 0x5555u;
    uint msbs = control_word & 0xaaaau;
    uint msbs_shift = msbs >> 1;
    msbs |= msbs_shift;
    return bitCount(lsbs) + bitCount(msbs) + q_bits * 8;
}

uint quantize_code_word(uint control_word, int quant)
{
    if (quant != 0 && control_word != 0)
    {
        int q_bits = int(bitfieldExtract(control_word, Q_PLANES_OFFSET, Q_PLANES_BITS));
        int sub_quant = min(q_bits, quant);
        q_bits -= sub_quant;
        quant -= sub_quant;

        if (quant != 0)
        {
            quant = min(quant, 3);

            uint plane0 = control_word & 0x5555u;
            uint plane1 = (control_word & 0xaaaau) >> 1;
            uint plane2 = plane0 & plane1;

            do
            {
                plane0 = plane1;
                plane1 = plane2;
                plane2 = 0;
                quant--;
            } while (quant != 0);

            plane0 &= ~plane1;

            uint new_control_word = plane0 | (plane1 << 1);
            control_word = bitfieldInsert(control_word, new_control_word, 0, 16);
        }

        control_word = bitfieldInsert(control_word, uint(q_bits), Q_PLANES_OFFSET, Q_PLANES_BITS);
    }

    return control_word;
}

uint copy_bytes(inout uint output_offset, uint input_offset, uint count)
{
    uint significant_mask = 0;

    do
    {
        uint in_data = uint(payload_data.data[input_offset]);
        // If we observe any 1 in the non-sign planes, it's not deadzone quantized.
        significant_mask |= in_data;
        bitstream_data_8b.data[output_offset++] = uint8_t(in_data);
        count--;
        input_offset++;
    } while (count > 0);

    return significant_mask;
}

uint modify_quant_code(uint code, int quant)
{
    int e = int(bitfieldExtract(code, 3, 5));
    e = max(e - quant, 0);
    code = bitfieldInsert(code, e, 3, 5);
    return code;
}

uint inclusive_add_clustered16(uint v)
{
    for (uint i = 1; i < 16; i *= 2)
    {
        uint up = subgroupShuffleUp(v, i);
        v += (gl_SubgroupInvocationID & 15) >= i ? up : 0;
    }

    return v;
}

shared uint shared_sign_bank[4][1024 / 32];
uint pending_sign_write = 0;
uint pending_sign_mask = 0;

void append_sign_plane(uint bank, inout uint local_sign_offset, uint sign_mask, uint significant_mask)
{
    // Clock out one bit a time. This seems kinda slow.
    while (significant_mask != 0)
    {
        int bit = findLSB(significant_mask);
        significant_mask &= significant_mask - 1;
        int out_bit = int(local_sign_offset & 31u);
        pending_sign_write = bitfieldInsert(pending_sign_write, bitfieldExtract(sign_mask, bit, 1), out_bit, 1);
        pending_sign_mask = bitfieldInsert(pending_sign_mask, 1, out_bit, 1);

        if (out_bit == 31)
        {
            if (pending_sign_mask == ~0u)
            {
                shared_sign_bank[bank][local_sign_offset / 32] = pending_sign_write;
            }
            else
            {
                atomicAnd(shared_sign_bank[bank][local_sign_offset / 32], ~pending_sign_mask);
                atomicOr(shared_sign_bank[bank][local_sign_offset / 32], pending_sign_write & pending_sign_mask);
            }

            pending_sign_mask = 0;
        }

        local_sign_offset++;
    }
}

void flush_sign_plane(uint bank, uint local_sign_offset)
{
    if (pending_sign_mask != 0)
    {
        atomicAnd(shared_sign_bank[bank][local_sign_offset / 32], ~pending_sign_mask);
        atomicOr(shared_sign_bank[bank][local_sign_offset / 32], pending_sign_write & pending_sign_mask);
        pending_sign_mask = 0;
    }
}

void main()
{
    uint index = gl_SubgroupInvocationID + gl_SubgroupSize * gl_SubgroupID;
    uint linear_block_32x32_index = index >> 4;
    ivec2 block32x32_index = 2 * ivec2(gl_WorkGroupID.xy);
    block32x32_index.x += int(bitfieldExtract(index, 4, 1));
    block32x32_index.y += int(bitfieldExtract(index, 5, 1));
    ivec2 local_block_index = ivec2(bitfieldExtract(index, 0, 2), bitfieldExtract(index, 2, 2));
    ivec2 block8x8_index = 4 * block32x32_index + local_block_index;

    BlockMeta meta;
    int quant;

    bool in_range_8x8 = all(lessThan(block8x8_index, registers.resolution_8x8_blocks));
    bool in_range_32x32 = all(lessThan(block32x32_index, registers.resolution_32x32_blocks));
    uint num_bits_for_q = 0;

    if (in_range_32x32)
    {
        int block_index = registers.block_offset_32x32 +
            registers.block_stride_32x32 * block32x32_index.y +
            block32x32_index.x;
        quant = quant_data.data[block_index];
    }
    else
    {
        quant = 0;
    }

    if (in_range_8x8)
    {
        int block_index = registers.block_offset_8x8 +
            registers.block_stride_8x8 * block8x8_index.y +
            block8x8_index.x;
        meta = block_meta.meta[block_index];
        uint num_planes = block_stats.stats[block_index].num_planes;
        num_bits_for_q = uint(block_stats.stats[block_index].errors[min(num_planes, quant)].payload_cost);
    }
    else
    {
        meta = BlockMeta(0, 0);
    }

    uint code_word = quantize_code_word(meta.code_word, quant);
    bool active_code_word = (code_word & 0xffffu) != 0;

    uvec4 code_word_ballot = subgroupBallot(active_code_word);
    uint local_ballot = gl_SubgroupSize >= 64 && linear_block_32x32_index >= 2 ? code_word_ballot.y : code_word_ballot.x;
    local_ballot = bitfieldExtract(local_ballot, int(16u * (linear_block_32x32_index & 1u)), 16);

    uint required_plane_bytes = compute_required_8x8_size(code_word);
    uint required_sign_bits = num_bits_for_q - required_plane_bytes * 8;

    uint required_bits_with_meta = num_bits_for_q;
    if (required_bits_with_meta != 0)
        required_bits_with_meta += 24;

    const uint HeaderSize = 2;
    bool writes_header =
        all(lessThan(block32x32_index, registers.resolution_32x32_blocks)) && (index & 15u) == 15u;

    uint payload_total_bits = subgroupClusteredAdd(required_bits_with_meta, 16);
    uint payload_total_words = (payload_total_bits + 31) / 32;
    if (payload_total_words != 0)
        payload_total_words += HeaderSize;

    uint global_payload_offset = 0;
    if (writes_header && payload_total_words != 0)
        global_payload_offset = atomicAdd(payload_data.bitstream_payload_counter, payload_total_words);
    global_payload_offset = subgroupShuffle(global_payload_offset, gl_SubgroupInvocationID | 15u);

    if (writes_header)
    {
        uint block_index = registers.block_offset_32x32 +
            block32x32_index.y * registers.block_stride_32x32 + block32x32_index.x;

        if (payload_total_words != 0)
        {
            bitstream_data.data[global_payload_offset + 0] =
                local_ballot | (payload_total_words << 16) | (registers.sequence_code << 28);
            bitstream_data.data[global_payload_offset + 1] =
                modify_quant_code(registers.quant_resolution_code, quant) | (block_index << 8);
        }

        bitstream_meta.packets[block_index] = BitstreamPacket(global_payload_offset, payload_total_words);
    }

    uint total_subblocks = bitCount(local_ballot);

    uint total_sign_bits = inclusive_add_clustered16(required_sign_bits);
    uint local_planes_offset = inclusive_add_clustered16(required_plane_bytes) - required_plane_bytes;
    uint local_sign_offset = total_sign_bits - required_sign_bits;
    uint global_planes_offset = 4 * global_payload_offset + 3 * total_subblocks + 4 * HeaderSize;
    uint global_sign_offset = global_planes_offset + subgroupClusteredAdd(required_plane_bytes, 16);
    global_planes_offset += local_planes_offset;

    uint total_sign_bytes = (subgroupShuffle(total_sign_bits, gl_SubgroupInvocationID | 15u) + 7) / 8;

    // Followed by N code words which map to the local ballot of active 16x16 regions.
    if (active_code_word)
    {
        uint block_header_offset = bitCount(bitfieldExtract(
            local_ballot, 0, local_block_index.y * 4 + local_block_index.x));

        uint in_q_bits = bitfieldExtract(meta.code_word, Q_PLANES_OFFSET, Q_PLANES_BITS);
        uint out_q_bits = bitfieldExtract(code_word, Q_PLANES_OFFSET, Q_PLANES_BITS);
        uint input_offset = meta.offset;
        uint output_offset = global_planes_offset;

        for (int bit_offset = 0; bit_offset < 16; bit_offset += 2)
        {
            uint out_planes = bitfieldExtract(code_word, bit_offset, 2) + out_q_bits;
            uint in_planes = bitfieldExtract(meta.code_word, bit_offset, 2) + in_q_bits;
            if (in_planes != 0)
                in_planes++;

            uint sign_plane = uint(payload_data.data[input_offset]);

            if (out_planes != 0)
            {
                uint significant_mask = copy_bytes(output_offset, input_offset + 1, out_planes);
                append_sign_plane(linear_block_32x32_index, local_sign_offset, sign_plane, significant_mask);
            }

            input_offset += in_planes;
        }

        flush_sign_plane(linear_block_32x32_index, local_sign_offset);

        bitstream_data_16b.data[2 * global_payload_offset + block_header_offset + 2 * HeaderSize] =
            uint16_t(code_word);
        bitstream_data_8b.data[4 * global_payload_offset + 2 * total_subblocks + block_header_offset + 4 * HeaderSize] =
            uint8_t(code_word >> 16);
    }

    subgroupBarrier();

    // Copy out all sign planes for any given group.
    for (uint i = index & 15u; i < total_sign_bytes / 4; i += 16)
    {
        uint sign_word = shared_sign_bank[linear_block_32x32_index][i];
        uint offset_8b = global_sign_offset + 4 * i;
        bitstream_data_8b.data[offset_8b + 0] = uint8_t(sign_word >> 0);
        bitstream_data_8b.data[offset_8b + 1] = uint8_t(sign_word >> 8);
        bitstream_data_8b.data[offset_8b + 2] = uint8_t(sign_word >> 16);
        bitstream_data_8b.data[offset_8b + 3] = uint8_t(sign_word >> 24);
    }

    // Copy out any stragglers.
    for (uint i = (total_sign_bytes & ~3u) + (index & 15u); i < total_sign_bytes; i += 16)
    {
        uint sign_word = shared_sign_bank[linear_block_32x32_index][i / 4];
        uint offset_8b = global_sign_offset + i;
        bitstream_data_8b.data[offset_8b] = uint8_t(sign_word >> (8 * (i & 3u)));
    }
}