debug_views.slang

#pragma once
#include "shared.h"
#include "lib/shadow_sample.slang"

#ifdef DEBUG

vec3 TurboColormap(f32 x) {
    // Source: https://research.google/blog/turbo-an-improved-rainbow-colormap-for-visualization/

    vec4 kRedVec4 = vec4(0.13572138, 4.61539260, -42.66032258, 132.13108234);
    vec4 kGreenVec4 = vec4(0.09140261, 2.19418839, 4.84296658, -14.18503333);
    vec4 kBlueVec4 = vec4(0.10667330, 12.64194608, -60.58204836, 110.36276771);
    vec2 kRedVec2 = vec2(-152.94239396, 59.28637943);
    vec2 kGreenVec2 = vec2(4.27729857, 2.82956604);
    vec2 kBlueVec2 = vec2(-89.90310912, 27.34824973);

    x = clamp(x, 0, 1);
    vec4 v4 = vec4(1.0, x, x * x, x * x * x);
    vec2 v2 = v4.zw * v4.z;
    return vec3(
        dot(v4, kRedVec4) + dot(v2, kRedVec2),
        dot(v4, kGreenVec4) + dot(v2, kGreenVec2),
        dot(v4, kBlueVec4) + dot(v2, kBlueVec2)
    );
}

bool try_debug_view_pre(
    VtxOut input,
    FrameUBO frame,
    inout vec4 albedo,
    inout f32 alpha,
    inout vec3 normalWS,
    inout f32 metallic,
    inout f32 roughness,
    inout vec3 emissive,
    vec3 N,
    vec3 T,
    vec3 B,
    vec3 normalTS,
    out vec4 result
) {

    DebugViewMode dbg_flag = frame.debug;
    f32 near_plane = frame.near_plane;
    f32 far_plane = frame.far_plane;

    switch (dbg_flag) {
    case DebugViewMode::DEPTH_grayscale:
        result = vec4(vec3(input.view_pos.z), 1.0);
        return true;

    case DebugViewMode::DEPTH_color:
        result = vec4(TurboColormap(abs(input.view_pos.z)), 1.0);
        return true;

    case DebugViewMode::UV:
        result = vec4(input.uv, 0.0, 1.0);
        return true;

    case DebugViewMode::NORMAL_TEX:
        result = vec4(normalTS * 0.5 + 0.5, 1.0);
        return true;

    case DebugViewMode::NORMAL_GEO:
        result = vec4(N * 0.5 + 0.5, 1.0);
        return true;

    case DebugViewMode::TANGENT_GEO:
        result = vec4(T * 0.5 + 0.5, 1.0);
        return true;

    case DebugViewMode::BITANGENT_GEO:
        result = vec4(B * 0.5 + 0.5, 1.0);
        return true;

    case DebugViewMode::TANGENT_GEO_W:
        result = vec4(vec3(input.tangent.w * 0.5 + 0.5), 1.0);
        return true;

    case DebugViewMode::SHADING_NORMAL:
        result = vec4(normalWS * 0.5 + 0.5, 1.0);
        return true;

    case DebugViewMode::ALPHA:
        result = vec4(vec3(alpha), 1.0);
        return true;

    case DebugViewMode::EMISSIVE:
        result = vec4(emissive, 1.0);
        return true;

    case DebugViewMode::ALBEDO:
        result = vec4(pow(albedo.rgb, 1.0 / 2.0), 1.0);
        return true;

    case DebugViewMode::METALLIC:
        result = vec4(vec3(metallic), 1.0);
        return true;

    case DebugViewMode::ROUGHNESS:
        result = vec4(vec3(roughness), 1.0);
        return true;

    case DebugViewMode::OCCLUSION:
        result = vec4(1.0, 0.0, 1.0, 1.0);
        return true;

    case DebugViewMode::SHADOW_CASCADE: {
        u32 cascade_idx = select_cascade_aabb(input.world_pos, frame.position_ws.xyz, frame.shadows);
        vec3 colors[3] = { vec3(1, 0, 0), vec3(0, 1, 0), vec3(0, 0, 1) };
        vec4 albed = vec4(pow(albedo.rgb, 1.0 / 2.0), 1.0);
        result = vec4(lerp(colors[cascade_idx], albed.xyz, 0.1), 0.25);
        return true;
    }

    case DebugViewMode::LIGHTING_ONLY: {
        albedo = vec4(0.5, 0.5, 0.5, 1.0);
        // roughness = 0.25;
        // metallic  = 0.0;
        // emissive  = 0.0;
        alpha = 1.0;
        // normalWS  = input.normal;
        return false;
    }

    default:
        return false;
    }
}

void apply_debug_view_post(inout vec3 color, VtxOut input, ClusterRecord cr, FrameUBO frame) {
    ClusterGrid clusters = frame.cluster;
    DebugViewMode dbg_flag = frame.debug;
    f32 near_plane = frame.near_plane;
    f32 far_plane = frame.far_plane;

    switch (dbg_flag) {
    case DebugViewMode::CLUSTER_HEATMAP: {
        f32 h = f32(cr.lightCount) / 16.0;
        vec4 heatmap = vec4(h, 0.5 - h, 0, 0.0);
        color = lerp(color.rgb, heatmap.rgb, 0.65);
        break;
    }

    case DebugViewMode::CLUSTER_BOUNDS: {
        uint tileX = uint(input.clip_pos.x) / TILE_SIZE_X;
        uint tileY = uint(input.clip_pos.y) / TILE_SIZE_Y;
        f32 vDepth = clamp(abs(input.view_pos.z), near_plane, far_plane);

        u32 tz = compute_depth_slice(vDepth, near_plane, far_plane, clusters.cluster_count_z);
        u32 flatIdx = flatten_cluster_index(tileX, tileY, tz, clusters.cluster_count_x, clusters.cluster_count_y);

        ClusterBounds c;
        c.minPoint = clusters.clusters_bounds[flatIdx].minPoint;
        c.maxPoint = clusters.clusters_bounds[flatIdx].maxPoint;

        bool tileBorder = (uint(input.clip_pos.x) % TILE_SIZE_X) < 2 || (uint(input.clip_pos.y) % TILE_SIZE_Y) < 2;
        f32 depthFrac = log(vDepth / near_plane) / log(far_plane / near_plane);
        f32 sliceCoord = depthFrac * f32(clusters.cluster_count_z);

        bool depthBorder = frac(sliceCoord) < 0.05;

        if (tileBorder) {
            color = lerp(color, vec3(1, 1, 1), 0.05);
        }

        if (depthBorder) {
            color = lerp(color, vec3(1, 0, 0), 0.5);
        }
        break;
    }

    default:
        break;
    }
}

#endif