ConstantBuffer<K_FloodSig> sig : register(b0);

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

Struct(CsInput)
{
    Semantic(uint3, SV_DispatchThreadID);
};

[numthreads(8, 8, 1)]
K_ENTRY void cs(CsInput input)
{
    uint2 id = input.SV_DispatchThreadID.xy;
    uint2 tex_size = uint2(sig.tex_width, sig.tex_height);
    if (id.x < tex_size.x && id.y < tex_size.y) {
        Texture2D<Vec4> emittance_tex = GpuResourceFromUrid(sig.emittance_tex_urid);
        RWTexture2D<uint2> read_flood_tex = GpuResourceFromUrid(sig.read_flood_tex_urid);
        RWTexture2D<uint2> target_flood_tex = GpuResourceFromUrid(sig.target_flood_tex_urid);
        int step_len = sig.step_len;
        if (step_len == -1) {
            /* Seed */
            Vec4 emittance = emittance_tex[id];
            uint2 seed = uint2(0xFFFF, 0xFFFF);
            if (emittance.a > 0) {
                seed = id;
            }
            target_flood_tex[id] = seed;
        } else {
            /* Flood */
            Vec2I32 read_coords[9] = {
                (Vec2I32)id + Vec2I32(-step_len,  -step_len),  /* top left */
                (Vec2I32)id + Vec2I32(0        ,  -step_len),  /* top center */
                (Vec2I32)id + Vec2I32(+step_len,  -step_len),  /* top right */
                (Vec2I32)id + Vec2I32(-step_len,  0        ),  /* center left */
                (Vec2I32)id + Vec2I32(0        ,  0        ),  /* center center */
                (Vec2I32)id + Vec2I32(+step_len,  0        ),  /* center right */
                (Vec2I32)id + Vec2I32(-step_len,  +step_len),  /* bottom left */
                (Vec2I32)id + Vec2I32(0        ,  +step_len),  /* bottom center */
                (Vec2I32)id + Vec2I32(+step_len,  +step_len)   /* bottom right */
            };
            uint2 closest_seed = uint2(0xFFFF, 0xFFFF);
            u32 closest_seed_len_sq = 0xFFFFFFFF;
            for (int i = 0; i < 9; ++i) {
                Vec2I32 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];
                    Vec2I32 dist_vec = (Vec2I32)id - (Vec2I32)seed;
                    u32 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;
        }
    }
}