#include "sh/common.hlsl"

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

#define ROOTSIG \
    "RootConstants(num32BitConstants = 16, b0), " \
    "DescriptorTable(SRV(t0, space = 0, numDescriptors = unbounded, flags = DESCRIPTORS_VOLATILE)), " \
    "DescriptorTable(SRV(t0, space = 1, numDescriptors = unbounded, flags = DESCRIPTORS_VOLATILE)), " \
    "DescriptorTable(SRV(t0, space = 2, numDescriptors = unbounded, flags = DESCRIPTORS_VOLATILE)), " \
    "DescriptorTable(UAV(u0, space = 3, numDescriptors = unbounded, flags = DESCRIPTORS_VOLATILE)), "

ConstantBuffer<struct sh_shade_constants> g_constants : register(b0);
Texture2D<float4> g_textures_float4[] : register(t0, space0);
Texture2D<uint2> g_textures_uint2[] : register(t0, space1);
Texture3D<uint> g_noise_textures[] : register(t0, space2);
RWTexture2D<float4> g_write_textures[]: register(u0, space3);

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

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

#define SAMPLES 4
#define MARCHES 16
#define AMBIENT float4(0, 0, 0, 0)

float rand_angle(int2 pos, uint ray_index) {
    Texture3D<uint> noise_tex = g_noise_textures[SH_BLUE_NOISE_TEX_ID];

    uint3 noise_coord = uint3(abs(pos.xy), ray_index);
    noise_coord += g_constants.frame_seed;
    // noise_coord.xy += g_constants.frame_seed.xy;

    uint noise = noise_tex[noise_coord % uint3(SH_BLUE_NOISE_TEX_WIDTH, SH_BLUE_NOISE_TEX_HEIGHT, SH_BLUE_NOISE_TEX_DEPTH)];
    return ((float)noise / (float)0xFFFF) * TAU;
}

INLINE float4 get_light_in_dir(uint2 ray_start, float2 ray_dir)
{
    Texture2D<uint2> flood_tex = g_textures_uint2[g_constants.emittance_flood_tex_urid];
    Texture2D<float4> emittance_tex = g_textures_float4[g_constants.emittance_tex_urid];

    float4 result = AMBIENT;
    float2 at_float = ray_start;
    uint2 at_uint = ray_start;
    for (uint i = 0; i < MARCHES; ++i) {
        uint2 flood = flood_tex[at_uint];
        float2 dist_vec = at_float - (float2)flood;
        float dist = length(dist_vec);
        if (dist < 1) {
            result = emittance_tex[flood];
            break;
        } else {
            at_float += ray_dir * dist;
            at_uint = round(at_float);
            if (at_uint.x < 0 || at_uint.x >= g_constants.tex_width || at_uint.y < 0 || at_uint.y >= g_constants.tex_height) {
                /* Ambient lighting (ray hit edge of screen) */
                break;
            }
        }
    }

    return result;
}

INLINE float4 get_light_at_pos(int2 pos)
{
    float4 result = 0;
    for (uint i = 0; i < SAMPLES; ++i) {
        float angle = rand_angle(pos, i);
        float2 dir = float2(cos(angle), sin(angle));
        float4 light_in_dir = get_light_in_dir(pos, dir);
        result += light_in_dir;
    }
    result /= SAMPLES;

    return result;
}

/* ========================== *
 * Tone map
 * ========================== */

/* ACES approximation by Krzysztof Narkowicz 
 * https://knarkowicz.wordpress.com/2016/01/06/aces-filmic-tone-mapping-curve/ */
INLINE float3 tone_map(float3 v)
{
    return saturate((v * (2.51f * v + 0.03f)) / (v * (2.43f * v + 0.59f) + 0.14f));
}

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

[numthreads(8, 8, 1)]
SH_ENTRY(ROOTSIG) void cs(struct cs_input input)
{
    uint2 id = input.SV_DispatchThreadID.xy;
    if (id.x < g_constants.tex_width && id.y < g_constants.tex_height) {
        Texture2D<float4> albedo_tex = g_textures_float4[g_constants.albedo_tex_urid];
        RWTexture2D<float4> write_tex = g_write_textures[g_constants.write_tex_urid];
        float4 color = float4(1, 1, 1, 1);

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

        /* Apply lighting */
        color *= get_light_at_pos(id);

        /* Apply tone map */
        /* TODO: Dynamic exposure based on average scene luminance */
        color.rgb = tone_map(color.rgb) * g_constants.exposure;

        /* Apply gamma correction */
        color = pow(abs(color), 1/g_constants.gamma);

        write_tex[id] = color;
    }
}