#include "kernel.h"

ConstantBuffer<K_ShadeSig> sig : register(b0);

/* ========================== *
 * Lighting
 * ========================== */

#define SAMPLES 16
#define MARCHES 16
#define EDGE_FALLOFF 100

float rand_angle(uint2 pos, u32 ray_index) {
    Texture3D<u32> noise_tex = GpuResourceFromUrid(K_BLUE_NOISE_TEX_ID);

    Vec3I32 noise_coord = Vec3I32(1, 1, 1);
    noise_coord += Vec3I32(pos.xy,  ray_index);
    noise_coord.xyz += sig.frame_seed.xyz;
    // noise_coord.xy -= sig.camera_offset;

    u32 noise = noise_tex[noise_coord % uint3(K_BLUE_NOISE_TEX_WIDTH, K_BLUE_NOISE_TEX_HEIGHT, K_BLUE_NOISE_TEX_DEPTH)];
    return ((float)noise / (float)0xFFFF) * Tau;
}

Vec3 get_light_in_dir(uint2 ray_start, Vec2 ray_dir)
{
    Texture2D<uint2> flood_tex = GpuResourceFromUrid(sig.emittance_flood_tex_urid);
    Texture2D<Vec4> emittance_tex = GpuResourceFromUrid(sig.emittance_tex_urid);

    Vec3 result = Vec3(0, 0, 0);
    Vec2 at_float = ray_start;
    uint2 at_uint = ray_start;
    for (u32 i = 0; i < MARCHES; ++i) {
        uint2 flood = flood_tex[at_uint];
        Vec2 dist_vec = at_float - (Vec2)flood;
        float dist = length(dist_vec);
        if (dist < 1) {
            /* Scale light by distance from edge so that offscreen-lights fade in/out rather than popping in */
            float dist_x = min(abs(sig.tex_width - at_float.x), at_float.x);
            float dist_y = min(abs(sig.tex_height - at_float.y), at_float.y);
            float dist_scale = min(min(dist_x, dist_y) / EDGE_FALLOFF, 1);
            result = emittance_tex[flood].rgb * dist_scale;
            break;
        } else {
            at_float += ray_dir * dist;
            at_uint = round(at_float);
            if (at_uint.x < 0 || at_uint.x >= sig.tex_width || at_uint.y < 0 || at_uint.y >= sig.tex_height) {
                /* Ray hit edge of screen */
                break;
            }
        }
    }
    return result;
}

Vec3 get_light_at_pos(uint2 pos)
{
    Vec3 result = 0;
    for (u32 i = 0; i < SAMPLES; ++i) {
        float angle = rand_angle(pos, i);
        Vec2 dir = Vec2(cos(angle), sin(angle));
        Vec3 light_in_dir = get_light_in_dir(pos, dir);
        result += light_in_dir;
    }
    result /= SAMPLES;
    return result;
}

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

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

[numthreads(8, 8, 1)]
K_ENTRY void cs(CsInput input)
{
    uint2 id = input.SV_DispatchThreadID.xy;
    if (id.x < sig.tex_width && id.y < sig.tex_height) {
        Texture2D<Vec4> albedo_tex = GpuResourceFromUrid(sig.albedo_tex_urid);
        Texture2D<Vec4> read_tex = GpuResourceFromUrid(sig.read_tex_urid);
        RWTexture2D<Vec4> target_tex = GpuResourceFromUrid(sig.target_tex_urid);
        Vec4 color = Vec4(1, 1, 1, 1);

        /* Apply albedo */
        color *= albedo_tex[id];

        /* Apply lighting */
        if (!(sig.flags & K_SHADE_FLAG_DISABLE_EFFECTS)) {
            color.rgb *= get_light_at_pos(id);
        }

        /* Apply temporal accumulation */
        float hysterisis = 0;
        // hysterisis = 0.2;
        // hysterisis = 0.4;
        // hysterisis = 0.5;
        // hysterisis = 0.9;
        color.rgb = lerp(color.rgb, read_tex[id].rgb, hysterisis);

        target_tex[id] = color;
    }
}