ConstantBuffer<MaterialSig>  mat_sig       : register (b0);
ConstantBuffer<FloodSig>     flood_sig     : register (b0);
ConstantBuffer<ShadeSig>     shade_sig     : register (b0);
ConstantBuffer<UiBlitSig>    ui_blit_sig   : register (b0);
ConstantBuffer<UiRectSig>    ui_rect_sig   : register (b0);
ConstantBuffer<UiShapeSig>   ui_shape_sig  : register (b0);

////////////////////////////////
//~ Material

Struct(MaterialPS_Input)
{
    Semantic(Vec4, sv_position);
    Semantic(nointerpolation Texture2DRid, tex);
    Semantic(nointerpolation u32, grid_id);
    Semantic(Vec2, uv);
    Semantic(Vec4, tint_lin);
    Semantic(Vec4, emittance_lin);
};

Struct(MaterialPS_Output)
{
    Semantic(Vec4, sv_target0);  /* Albedo */
    Semantic(Vec4, sv_target1);  /* Emittance */
};

//- Vertex shader

MaterialPS_Input VSDef(MaterialVS, Semantic(u32, sv_instanceid), Semantic(u32, sv_vertexid))
{
    ConstantBuffer<MaterialSig> sig = mat_sig;
    static const Vec2 unit_quad_verts[4] = {
        Vec2(-0.5f, -0.5f),
        Vec2(0.5f, -0.5f),
        Vec2(0.5f, 0.5f),
        Vec2(-0.5f, 0.5f)
    };

    StructuredBuffer<MaterialInstance> instances = UniformResourceFromRid(sig.instances);
    Vec2 vert = unit_quad_verts[sv_vertexid];
    MaterialInstance instance = instances[sv_instanceid];

    Vec2 world_pos = mul(instance.xf, Vec3(vert, 1)).xy;

    MaterialPS_Input result;
    result.sv_position = mul(sig.projection, Vec4(world_pos, 0, 1));
    result.tex = instance.tex;
    result.grid_id = instance.grid_id;
    result.uv = instance.uv0 + ((vert + 0.5) * (instance.uv1 - instance.uv0));
    result.tint_lin = LinearFromSrgbU32(instance.tint_srgb);
    result.emittance_lin = LinearFromSrgbVec4(Vec4(instance.light_emittance_srgb, instance.is_light));
    return result;
}

//- Pixel shader

MaterialPS_Output PSDef(MaterialPS, MaterialPS_Input input)
{
    ConstantBuffer<MaterialSig> sig = mat_sig;

    MaterialPS_Output result;
    Vec4 albedo = input.tint_lin;

    /* Texture */
    if (input.tex < 0xFFFFFFFF)
    {
        SamplerState sampler = UniformSamplerFromRid(sig.sampler);
        Texture2D<Vec4> tex = NonUniformResourceFromRid(input.tex);
        albedo *= tex.Sample(sampler, input.uv);
    }

    /* Grid */
    if (input.grid_id < 0xFFFFFFFF)
    {
        // StructuredBuffer<MaterialGrid> grids = UniformResourceFromRid(sig.grids);
        StructuredBuffer<MaterialGrid> grids = ResourceDescriptorHeap[sig.grids];
        MaterialGrid grid = grids[input.grid_id];
        Vec2 grid_pos = input.sv_position.xy + grid.offset;
        f32 half_thickness = grid.line_thickness / 2;
        f32 spacing = grid.line_spacing;
        u32 color_srgb = grid.bg0_srgb;
        Vec2 v = abs(round(grid_pos / spacing) * spacing - grid_pos);
        f32 dist = min(v.x, v.y);
        if (grid_pos.y <= half_thickness && grid_pos.y >= -half_thickness)
        {
            color_srgb = grid.x_srgb;
        }
        else if (grid_pos.x <= half_thickness && grid_pos.x >= -half_thickness)
        {
            color_srgb = grid.y_srgb;
        }
        else if (dist < half_thickness)
        {
            color_srgb = grid.line_srgb;
        }
        else
        {
            bool checker = 0;
            u32 cell_x = (u32)(abs(grid_pos.x) / spacing) + (grid_pos.x < 0);
            u32 cell_y = (u32)(abs(grid_pos.y) / spacing) + (grid_pos.y < 0);
            if (cell_x % 2 == 0)
            {
                checker = cell_y % 2 == 0;
            }
            else
            {
                checker = cell_y % 2 == 1;
            }
            if (checker)
            {
                color_srgb = grid.bg1_srgb;
            }
        }
        albedo = LinearFromSrgbU32(color_srgb);
    }

    Vec4 emittance = input.emittance_lin * albedo.a;

    result.sv_target0 = albedo;
    result.sv_target1 = emittance;

    return result;
}

////////////////////////////////
//~ Flood

//- Compute shader

[numthreads(8, 8, 1)]
void CSDef(FloodCS, Semantic(Vec3U32, sv_dispatchthreadid))
{
    ConstantBuffer<FloodSig> sig = flood_sig;

    Vec2U32 id = sv_dispatchthreadid.xy;
    Vec2U32 tex_size = Vec2U32(sig.tex_width, sig.tex_height);
    if (id.x < tex_size.x && id.y < tex_size.y)
    {
        Texture2D<Vec4> emittance_tex = UniformResourceFromRid(sig.emittance);
        RWTexture2D<Vec2U32> read_flood_tex = UniformResourceFromRid(sig.read);
        RWTexture2D<Vec2U32> target_flood_tex = UniformResourceFromRid(sig.target);
        i32 step_len = sig.step_len;
        if (step_len == -1)
        {
            /* Seed */
            Vec4 emittance = emittance_tex[id];
            Vec2U32 seed = Vec2U32(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 */
            };
            Vec2U32 closest_seed = Vec2U32(0xFFFF, 0xFFFF);
            u32 closest_seed_len_sq = 0xFFFFFFFF;
            for (i32 i = 0; i < 9; ++i)
            {
                Vec2I32 coord = read_coords[i];
                if (coord.x >= 0 && coord.x < (i32)tex_size.x && coord.y >= 0 && coord.y < (i32)tex_size.y)
                {
                    Vec2U32 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;
        }
    }
}

////////////////////////////////
//~ Shade

#define LightSamples 16
#define LightMarches 16
#define LightEdgeFalloff 100

f32 RandAngle(Vec2U32 pos, u32 ray_index)
{
    ConstantBuffer<ShadeSig> sig = shade_sig;
    Texture3D<u32> noise_tex = UniformResourceFromRid(sig.noise);

    Vec3I32 noise_coord = Vec3U32(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 % Vec3U32(sig.noise_tex_width, sig.noise_tex_height, sig.noise_tex_depth)];
    return ((f32)noise / (f32)0xFFFF) * Tau;
}

Vec3 ColorFromDir(Vec2U32 ray_start, Vec2 ray_dir)
{
    ConstantBuffer<ShadeSig> sig = shade_sig;
    Texture2D<Vec2U32> flood_tex = UniformResourceFromRid(sig.emittance_flood);
    Texture2D<Vec4> emittance_tex = UniformResourceFromRid(sig.emittance);

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

Vec3 ColorFromPos(Vec2U32 pos)
{
    Vec3 result = 0;
    for (u32 i = 0; i < LightSamples; ++i)
    {
        f32 angle = RandAngle(pos, i);
        Vec2 dir = Vec2(cos(angle), sin(angle));
        Vec3 light_in_dir = ColorFromDir(pos, dir);
        result += light_in_dir;
    }
    result /= LightSamples;
    return result;
}

//- Compute shader

[numthreads(8, 8, 1)]
void CSDef(ShadeCS, Semantic(Vec3U32, sv_dispatchthreadid))
{
    ConstantBuffer<ShadeSig> sig = shade_sig;

    Vec2U32 id = sv_dispatchthreadid.xy;
    if (id.x < sig.tex_width && id.y < sig.tex_height)
    {
        Texture2D<Vec4> albedo_tex = UniformResourceFromRid(sig.albedo);
        Texture2D<Vec4> read_tex = UniformResourceFromRid(sig.read);
        RWTexture2D<Vec4> target_tex = UniformResourceFromRid(sig.target);
        Vec4 color = Vec4(1, 1, 1, 1);

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

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

        /* Apply temporal accumulation */
        f32 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;
    }
}

////////////////////////////////
//~ Ui Blit

Struct(UiBlitPS_Input)
{
    Semantic(Vec4, sv_position);
    Semantic(Vec2, uv);
};

Struct(UiBlitPS_Output)
{
    Semantic(Vec4, sv_target);
};

//- ACES

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

//- Vertex shader

UiBlitPS_Input VSDef(UiBlitVS, Semantic(u32, sv_vertexid))
{
    ConstantBuffer<UiBlitSig> sig = ui_blit_sig;
    static const Vec2 unit_quad_verts[4] = {
        Vec2(-0.5f, -0.5f),
        Vec2(0.5f, -0.5f),
        Vec2(0.5f, 0.5f),
        Vec2(-0.5f, 0.5f)
    };
    Vec2 vert = unit_quad_verts[sv_vertexid];

    UiBlitPS_Input result;
    result.sv_position = mul(sig.projection, Vec4(vert, 0, 1));
    result.uv = vert + 0.5;
    return result;
}

//- Pixel shader

UiBlitPS_Output PSDef(UiBlitPS, UiBlitPS_Input input)
{
    ConstantBuffer<UiBlitSig> sig = ui_blit_sig;
    SamplerState sampler = UniformSamplerFromRid(sig.sampler);

    UiBlitPS_Output result;
    Texture2D<Vec4> tex = UniformResourceFromRid(sig.src);
    Vec4 color = tex.Sample(sampler, input.uv);

    /* Apply tone map */
    if (sig.flags & UiBlitFlag_ToneMap)
    {
        /* TODO: Dynamic exposure based on average scene luminance */
        color.rgb *= sig.exposure;
        color.rgb = ToneMap(color.rgb);
    }

    /* Apply gamma correction */
    if (sig.flags & UiBlitFlag_GammaCorrect)
    {
        color = pow(abs(color), 1/sig.gamma);
    }

    result.sv_target = color;
    return result;
}

////////////////////////////////
//~ Ui rect

Struct(UiRectPS_Input)
{
    Semantic(Vec4, sv_position);
    Semantic(nointerpolation Texture2DRid, tex);
    Semantic(Vec2, uv);
    Semantic(Vec4, tint_srgb);
};

Struct(UiRectPS_Output)
{
    Semantic(Vec4, sv_target0);
};

//- Vertex shader

UiRectPS_Input VSDef(UiRectVS, Semantic(u32, sv_instanceid), Semantic(u32, sv_vertexid))
{
    ConstantBuffer<UiRectSig> sig = ui_rect_sig;
    static const Vec2 unit_quad_verts[4] = {
        Vec2(-0.5f, -0.5f),
        Vec2(0.5f, -0.5f),
        Vec2(0.5f, 0.5f),
        Vec2(-0.5f, 0.5f)
    };

    StructuredBuffer<UiRectInstance> instances = UniformResourceFromRid(sig.instances);
    UiRectInstance instance = instances[sv_instanceid];
    Vec2 vert = unit_quad_verts[sv_vertexid];
    Vec2 world_pos = mul(instance.xf, Vec3(vert, 1)).xy;

    UiRectPS_Input result;
    result.sv_position = mul(sig.projection, Vec4(world_pos, 0, 1));
    result.tex = instance.tex;
    result.uv = instance.uv0 + ((vert + 0.5) * (instance.uv1 - instance.uv0));
    result.tint_srgb = Vec4NormFromU32(instance.tint_srgb);
    return result;
}

//- Pixel shader

UiRectPS_Output PSDef(UiRectPS, UiRectPS_Input input)
{
    ConstantBuffer<UiRectSig> sig = ui_rect_sig;
    UiRectPS_Output result;
    Vec4 color = input.tint_srgb;

    /* Texture */
    if (input.tex < 0xFFFFFFFF)
    {
        Texture2D<Vec4> tex = NonUniformResourceFromRid(input.tex);
        SamplerState sampler = UniformSamplerFromRid(sig.sampler);
        color *= tex.Sample(sampler, input.uv);
    }

    result.sv_target0 = color;
    return result;
}

////////////////////////////////
//~ Ui shape

Struct(UiShapePS_Input)
{
    Semantic(Vec4, sv_position);
    Semantic(Vec4, color_srgb);
};

Struct(UiShapePS_Output)
{
    Semantic(Vec4, sv_target);
};

//- Vertex shader

UiShapePS_Input VSDef(UiShapeVS, Semantic(u32, sv_vertexid))
{
    ConstantBuffer<UiShapeSig> sig = ui_shape_sig;
    StructuredBuffer<UiShapeVert> verts = UniformResourceFromRid(sig.verts);
    UiShapeVert vert = verts[sv_vertexid];
    UiShapePS_Input result;
    result.sv_position = mul(sig.projection, Vec4(vert.pos.xy, 0, 1));
    result.color_srgb = Vec4NormFromU32(vert.color_srgb);
    return result;
}

//- Pixel shader

UiShapePS_Output PSDef(UiShapePS, UiShapePS_Input input)
{
    UiShapePS_Output result;
    result.sv_target = input.color_srgb;
    return result;
}