#include "common.hlsl"

/* ========================== *
 * Root signature
 * ========================== */

#define ROOTSIG \
    "RootConstants(num32BitConstants = 8, b0), " \
    "DescriptorTable(SRV(t0, space = 0, numDescriptors = unbounded, flags = DESCRIPTORS_VOLATILE)), " \
    "DescriptorTable(UAV(u0, space = 1, numDescriptors = unbounded, flags = DESCRIPTORS_VOLATILE)), " \
    \
    \
    "StaticSampler(s0, " \
                  "filter = FILTER_MIN_MAG_MIP_POINT, " \
                  "addressU = TEXTURE_ADDRESS_CLAMP, " \
                  "addressV = TEXTURE_ADDRESS_CLAMP, " \
                  "addressW = TEXTURE_ADDRESS_CLAMP, " \
                  "maxAnisotropy = 1)"

ConstantBuffer<struct sh_flood_constants> g_constants : register(b0);
Texture2D<float4> g_emittance_textures[] : register(t0, space0);
RWTexture2D<uint2> g_flood_textures[]: register(u0, space1);

SamplerState g_sampler : register(s0);

/* ========================== *
 * Entry point
 * ========================== */

struct cs_input {
    DECLS(uint3, SV_DispatchThreadID);
};

[numthreads(8, 8, 1)]
SH_ENTRY(ROOTSIG) void cs(struct cs_input input)
{
    uint2 id = input.SV_DispatchThreadID.xy;
    uint2 tex_size = uint2(g_constants.tex_width, g_constants.tex_height);
    if (id.x < tex_size.x && id.y < tex_size.y) {
        Texture2D<float4> emittance_tex = g_emittance_textures[g_constants.emittance_tex_urid];
        RWTexture2D<uint2> read_flood_tex = g_flood_textures[g_constants.read_flood_tex_urid];
        RWTexture2D<uint2> target_flood_tex = g_flood_textures[g_constants.target_flood_tex_urid];
        int step_len = g_constants.step_len;
        if (step_len == -1) {
            /* Seed */
            float4 emittance = emittance_tex[id];
            uint2 seed = uint2(0xFFFF, 0xFFFF);
            if (emittance.a > 0) {
                seed = id;
            }
            target_flood_tex[id] = seed;
        } else {
            /* Flood */
            int2 read_coords[9] = {
                (int2)id + int2(-step_len,  -step_len),  /* top left */
                (int2)id + int2(0        ,  -step_len),  /* top center */
                (int2)id + int2(+step_len,  -step_len),  /* top right */
                (int2)id + int2(-step_len,  0        ),  /* center left */
                (int2)id + int2(0        ,  0        ),  /* center center */
                (int2)id + int2(+step_len,  0        ),  /* center right */
                (int2)id + int2(-step_len,  +step_len),  /* bottom left */
                (int2)id + int2(0        ,  +step_len),  /* bottom center */
                (int2)id + int2(+step_len,  +step_len)   /* bottom right */
            };
            uint2 closest_seed = uint2(0xFFFF, 0xFFFF);
            uint closest_seed_len_sq = 0xFFFFFFFF;
            for (int i = 0; i < 9; ++i) {
                int2 coord = read_coords[i];
                if (coord.x >= 0 && coord.x < (int)tex_size.x && coord.y >= 0 && coord.y < (int)tex_size.y) {
                    uint2 seed = read_flood_tex[coord];
                    int2 dist_vec = (int2)id - (int2)seed;
                    uint dist_len_sq = dot(dist_vec, dist_vec);
                    if (dist_len_sq < closest_seed_len_sq) {
                        closest_seed = seed;
                        closest_seed_len_sq = dist_len_sq;
                    }
                }
            }
            target_flood_tex[id] = closest_seed;
        }
    }
}