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

#include "dwt_quant_scale.h"
#include "constants.h"

layout(local_size_x = 128) in;
layout(constant_id = 1) const bool SkipQuantScale = false;

layout(set = 0, binding = 0) uniform sampler2DArray uTexture;

struct QuantStats
{
    float16_t square_error;
    uint16_t payload_cost;
};

struct BlockMeta
{
    uint code_word;
    uint offset;
};

// Fit into 64 bytes.
struct BlockStats
{
    uint num_planes;
    QuantStats errors[15];
};

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

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

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

#include "dwt_swizzle.h"

layout(push_constant) uniform Registers
{
    ivec2 resolution;
    ivec2 resolution_8x8_blocks;
    vec2 inv_resolution;
    float input_layer;
    float quant_resolution;
    int block_offset;
    int block_stride;
    float rdo_distortion_scale;
} registers;

float max4(vec4 v)
{
    vec2 v2 = max(v.xy, v.zw);
    return max(v2.x, v2.y);
}

int max4(ivec4 v)
{
    ivec2 v2 = max(v.xy, v.zw);
    return max(v2.x, v2.y);
}

int scan_clustered8(int v)
{
    for (uint i = 1; i < 8; i *= 2)
    {
        int up = subgroupShuffleUp(v, i);
        v += (gl_SubgroupInvocationID & 7u) >= i ? up : 0;
    }

    return v;
}

void compute_quant_scale(float max_wave_texels, out uint quant_code, out float quant_scale)
{
    if (SkipQuantScale || max_wave_texels < 1.0)
    {
        quant_code = ENCODE_QUANT_IDENTITY;
        quant_scale = 1.0;
    }
    else
    {
        int e;
        frexp(max_wave_texels - 0.25, e);
        float target_max = float(1 << e) - 0.25;
        float inv_scale = max_wave_texels / target_max;
        quant_code = encode_quant_scale(inv_scale);
        quant_scale = 1.0 / decode_quant_scale(quant_code);
    }
}

float compute_square_error(mat2x4 v, int q, out uint num_significant_values)
{
    v = mat2x4(abs(v[0]), abs(v[1]));
    mat2x4 iv = mat2x4(floor(ldexp(v[0], ivec4(-q))), trunc(ldexp(v[1], ivec4(-q))));
    num_significant_values = 0;
    for (int j = 0; j < 2; j++)
        for (int i = 0; i < 4; i++)
            if (iv[j][i] != 0.0)
                num_significant_values++;
    iv[0] += mix(vec4(0.0), vec4(0.5), notEqual(iv[0], vec4(0.0)));
    iv[1] += mix(vec4(0.0), vec4(0.5), notEqual(iv[1], vec4(0.0)));
    iv = mat2x4(trunc(ldexp(iv[0], ivec4(q))), trunc(ldexp(iv[1], ivec4(q))));
    mat2x4 err = v - iv;
    num_significant_values = subgroupClusteredAdd(num_significant_values, 8);
    return (dot(err[0], err[0]) + dot(err[1], err[1])) * registers.rdo_distortion_scale;
}

struct QuantResult
{
    float square_error;
    int encode_cost_early;
    int block4x2_shifted;
    int encode_cost_late_bits;
    int quality_planes;
};

QuantResult compute_quant_stats(mat2x4 v, int q, int msb, int block4x2_max, float inv_quant_squared)
{
    block4x2_max >>= q;

    uint wave8_num_significants;
    QuantResult result;

    result.square_error = compute_square_error(v, q, wave8_num_significants) * inv_quant_squared;
    result.block4x2_shifted = block4x2_max;

    result.encode_cost_early = block4x2_max > 0 ? 1 : 0;
    msb -= q;

    result.quality_planes = 0;

    if (msb >= 3)
    {
        result.quality_planes = msb - 2;
        // Must encode the sign plane if we have quality planes.
        result.encode_cost_early = result.quality_planes + 1;
        result.block4x2_shifted >>= result.quality_planes;
    }

    result.encode_cost_early += findMSB(result.block4x2_shifted) + 1;
    result.encode_cost_late_bits = 8 * subgroupClusteredAdd(max(result.encode_cost_early - 1, 0), 8) + int(wave8_num_significants);
    return result;
}

float square(float v)
{
    return v * v;
}

void encode_payload(ivec2 block_index_8x8, mat2x4 texels)
{
    precise float max_subblock_texel = max(max4(abs(texels[0])), max4(abs(texels[1])));
    precise float max_wave_texels = subgroupClusteredMax(max_subblock_texel, 8);
    float quant_scale;
    uint quant_code;
    compute_quant_scale(max_wave_texels, quant_code, quant_scale);
    texels *= quant_scale;
    max_wave_texels *= quant_scale;
    max_subblock_texel *= quant_scale;

    float overall_quant_scale = registers.quant_resolution * quant_scale;
    float inv_quant = 1.0 / overall_quant_scale;
    float inv_quant_squared = inv_quant * inv_quant;
    ivec4 abs_quant_texels0 = abs(ivec4(texels[0]));
    ivec4 abs_quant_texels1 = abs(ivec4(texels[1]));
    int max_absolute_value = int(max_wave_texels);
    int block4x2_max = int(max_subblock_texel);

    uint block_index = registers.block_offset + block_index_8x8.y * registers.block_stride + block_index_8x8.x;

    // The entire block quantizes to zero.
    if (max_absolute_value == 0)
    {
        if ((gl_SubgroupInvocationID & 7) == 0)
        {
            block_meta.meta[block_index] = BlockMeta(0, 0);
            block_stats.stats[block_index].num_planes = 0;
            block_stats.stats[block_index].errors[0] = QuantStats(float16_t(0.0), uint16_t(0));
        }
        return;
    }

    int msb = findMSB(max_absolute_value);

    QuantResult result = compute_quant_stats(texels, 0, msb, block4x2_max, inv_quant_squared);
    int scan = scan_clustered8(result.encode_cost_early);

    uint global_offset = 0;

    // For feedback, and allocation of payload.
    if ((gl_SubgroupInvocationID & 7u) == 7u)
        global_offset = atomicAdd(payload_data.counter, scan);
    global_offset = subgroupShuffle(global_offset, gl_SubgroupInvocationID | 7u);

    scan -= result.encode_cost_early;

    // First, encode the code word.
    int quality_planes = result.quality_planes;
    uint code_word = quality_planes << Q_PLANES_OFFSET;
    code_word = bitfieldInsert(code_word, quant_code, QUANT_SCALE_OFFSET, QUANT_SCALE_BITS);
    uint plane_code = findMSB(result.block4x2_shifted) + 1;

    uint merged_plane_code = plane_code << ((gl_SubgroupInvocationID & 7u) * 2u);
    merged_plane_code |= subgroupShuffleXor(merged_plane_code, 1u);
    merged_plane_code |= subgroupShuffleXor(merged_plane_code, 2u);
    merged_plane_code |= subgroupShuffleXor(merged_plane_code, 4u);
    code_word |= merged_plane_code;

    if ((gl_SubgroupInvocationID & 7u) == 0u)
    {
        block_meta.meta[block_index] = BlockMeta(code_word, global_offset);
        block_stats.stats[block_index].num_planes = msb + 1;
        block_stats.stats[block_index].errors[0] = QuantStats(
            float16_t(0.0), // We don't care about distortion from 0 quant since we've already made that decision.
            uint16_t(result.encode_cost_late_bits));
    }

    for (int q = 1; q <= msb; q++)
    {
        QuantResult quant_result = compute_quant_stats(texels, q, msb, block4x2_max, inv_quant_squared);
        float square_error = subgroupClusteredAdd(quant_result.square_error, 8);

        if ((gl_SubgroupInvocationID & 7u) == 0)
        {
            block_stats.stats[block_index].errors[q] = QuantStats(
                float16_t(min(square_error, 60000.0)),
                uint16_t(quant_result.encode_cost_late_bits));
        }
    }

    // Record distortion for throwing away everything.
    float square_error = subgroupClusteredAdd((dot(texels[0], texels[0]) + dot(texels[1], texels[1])) * inv_quant_squared, 8);
    if ((gl_SubgroupInvocationID & 7u) == 0)
        block_stats.stats[block_index].errors[msb + 1] = QuantStats(float16_t(min(60000.0, square_error)), uint16_t(0));

    uint byte_offset = scan + global_offset;
    bool need_sign = result.block4x2_shifted != 0 || quality_planes != 0;

    // Don't pack the sign plane until final pass, since we don't know how we quantize yet.
    if (need_sign)
    {
        uvec4 s0 = uvec4(lessThan(texels[0], vec4(0.0))) << uvec4(0, 1, 2, 3);
        uvec4 s1 = uvec4(lessThan(texels[1], vec4(0.0))) << uvec4(4, 5, 6, 7);
        uint s = s0.x | s0.y | s0.z | s0.w | s1.x | s1.y | s1.z | s1.w;
        payload_data.data[byte_offset++] = uint8_t(s);

        int plane_iterations = quality_planes + int(plane_code);
        int q = plane_iterations - 1;
        do
        {
            s0 = uvec4(
                bitfieldExtract(uint(abs_quant_texels0.x), q, 1),
                bitfieldExtract(uint(abs_quant_texels0.y), q, 1),
                bitfieldExtract(uint(abs_quant_texels0.z), q, 1),
                bitfieldExtract(uint(abs_quant_texels0.w), q, 1));
            s1 = uvec4(
                bitfieldExtract(uint(abs_quant_texels1.x), q, 1),
                bitfieldExtract(uint(abs_quant_texels1.y), q, 1),
                bitfieldExtract(uint(abs_quant_texels1.z), q, 1),
                bitfieldExtract(uint(abs_quant_texels1.w), q, 1));
            s0 <<= uvec4(0, 1, 2, 3);
            s1 <<= uvec4(4, 5, 6, 7);
            s = s0.x | s0.y | s0.z | s0.w | s1.x | s1.y | s1.z | s1.w;
            payload_data.data[byte_offset++] = uint8_t(s);
            q--;
        } while (q >= 0);
    }
}

void main()
{
    uint local_index = gl_SubgroupID * gl_SubgroupSize + gl_SubgroupInvocationID;
    uint block_local_index = bitfieldExtract(local_index, 0, 3);
    uint block_x = bitfieldExtract(local_index, 3, 2);
    uint block_y = bitfieldExtract(local_index, 5, 2);

    // Each thread individually encodes 8 values.
    ivec2 local_coord = unswizzle8x8(block_local_index << 3);

    ivec2 coord = ivec2(gl_WorkGroupID.xy) * 32;
    coord += 8 * ivec2(block_x, block_y);
    coord += local_coord;

    ivec2 block_index = 4 * ivec2(gl_WorkGroupID.xy) + ivec2(block_x, block_y);

    vec3 uv = vec3(vec2(coord) * registers.inv_resolution, registers.input_layer);
    vec4 texels0 = textureGatherOffset(uTexture, uv, ivec2(1, 1)).wxzy;
    vec4 texels1 = textureGatherOffset(uTexture, uv, ivec2(3, 1)).wxzy;
    precise vec4 scaled_texels0 = texels0 * registers.quant_resolution;
    precise vec4 scaled_texels1 = texels1 * registers.quant_resolution;
    bool in_bounds = all(lessThan(block_index, registers.resolution_8x8_blocks));
    if (in_bounds)
        encode_payload(block_index, mat2x4(scaled_texels0, scaled_texels1));
}