@vs vs_main

layout(location = 0) in vec3 positionOS;
layout(location = 1) in vec2 texCoord;
layout(location = 2) in vec3 normalOS;
layout(location = 3) in vec4 tangentOS;

layout(location = 0) out vec3 var_positionVS;
layout(location = 1) out vec2 var_texCoord;
layout(location = 2) out vec3 var_normalVS;
layout(location = 3) out vec3 var_tangentVS;
layout(location = 4) out vec3 var_bitangentVS;

layout(binding = 0) uniform Global_Uniforms_Vertex {
    mat4 _MatrixWStoVS;
    mat4 _MatrixVStoCS;
};

layout(binding = 1) uniform Object_Uniforms {
    mat4 _MatrixOStoWS;
    mat4 _MatrixOStoWSNormal;
};

void main() {
    vec3 positionWS = (_MatrixOStoWS * vec4(positionOS, 1.0)).xyz;
    vec3 positionVS = (_MatrixWStoVS * vec4(positionWS, 1.0)).xyz;
    vec4 positionCS = _MatrixVStoCS * vec4(positionVS, 1.0);

    vec3 normalWS = normalize((_MatrixOStoWSNormal * vec4(normalOS, 0.0)).xyz);
    vec3 normalVS = normalize((_MatrixWStoVS * vec4(normalWS, 0.0)).xyz);

    vec3 tangentWS = normalize((_MatrixOStoWS * vec4(tangentOS.xyz, 0.0)).xyz);
    vec3 tangentVS = normalize((_MatrixWStoVS * vec4(tangentWS, 0.0)).xyz);

    vec3 bitangentVS = tangentOS.w * normalize(cross(normalVS, tangentVS));

    gl_Position = positionCS;
    var_positionVS = positionVS;
    var_texCoord = texCoord;
    var_normalVS = normalVS;
    var_tangentVS = tangentVS;
    var_bitangentVS = bitangentVS;
}
@end

@fs fs_main

struct Point_Light {
    vec3 positionWS;
    vec3 color;
};

struct Directional_Light {
    vec3 directionWS;
    vec3 color;
};

layout(location = 0) in vec3 var_positionVS;
layout(location = 1) in vec2 var_texCoord;
layout(location = 2) in vec3 var_normalVS;
layout(location = 3) in vec3 var_tangentVS;
layout(location = 4) in vec3 var_bitangentVS;

out vec4 fragColor;

layout(binding = 2) uniform Global_Uniforms_Fragment {
    mat4 _MatrixWStoVS;
    vec3 _AmbientLight;
    int _PointLightCount;
    int _DirectionalLightCount;
};

layout(binding = 3) uniform Material_Uniforms {
    int _TextureIndex;
};

layout(binding = 0) readonly buffer Point_Lights {
    Point_Light _PointLights[];
};

layout(binding = 1) readonly buffer Directional_Lights {
    Directional_Light _DirectionalLights[];
};

layout(binding = 2) uniform texture2DArray _BaseColorTexture;
layout(binding = 3) uniform texture2DArray _OcclusionRoughnessMetallicTexture;
layout(binding = 4) uniform texture2DArray _NormalTexture;

layout(binding = 0) uniform sampler _Sampler;

const float INV_PI = 0.31830987;
const float IOR = 1.45;
const vec3 DIELECTRIC_F0 = vec3(pow((IOR - 1.0) / (IOR + 1.0), 2.0));
const vec3 F90 = vec3(1.0);

vec3 fresnelSchlick(float dotVH, vec3 f0) {
    return mix(f0, F90, pow(1.0 - dotVH, 5.0));
}

float visibilityGGX(float dotNL, float dotNV, float alpha) {
    float alphaSquared = alpha * alpha;

    float vGGX = dotNL * sqrt(dotNV * dotNV * (1.0 - alphaSquared) + alphaSquared);
    float lGGX = dotNV * sqrt(dotNL * dotNL * (1.0 - alphaSquared) + alphaSquared);
    float GGX = vGGX + lGGX;
    return mix(0.0, 0.5 / GGX, GGX > 0.0);
}

float distributionGGX(float dotNH, float alpha) {
    float alphaSquared = alpha * alpha;
    float tmp = dotNH * dotNH * (alphaSquared - 1.0) + 1.0;
    return alphaSquared * INV_PI / (tmp * tmp);
}

vec3 toneMapAcesNarkowicz(vec3 color) {
    const float A = 2.51;
    const float B = 0.03;
    const float C = 2.43;
    const float D = 0.59;
    const float E = 0.14;
    return clamp((color * (A * color + B)) / (color * (C * color + D) + E), 0.0, 1.0);
}

vec3 lightOutgoingRadiance(
    vec3 viewDirectionVS, vec3 normalVS, float dotNV,
    vec3 baseColor, float alpha, float metallic, vec3 f0,
    vec3 incomingRadiance, vec3 lightDirectionVS
) {
    vec3 halfVectorVS = normalize(lightDirectionVS + viewDirectionVS);
    float dotVH = clamp(dot(viewDirectionVS, halfVectorVS), 0.0, 1.0);
    float dotNH = clamp(dot(normalVS, halfVectorVS), 0.0, 1.0);
    float dotNL = clamp(dot(normalVS, lightDirectionVS), 0.0, 1.0);

    vec3 fresnel = fresnelSchlick(dotVH, f0);
    float visibility = visibilityGGX(dotNL, dotNV, alpha);
    float distribution = distributionGGX(dotNH, alpha);

    vec3 scatteredFactor = (1.0 - fresnel) * (1.0 - metallic) * baseColor * INV_PI;
    vec3 reflectedFactor = fresnel * visibility * distribution;

    return (scatteredFactor + reflectedFactor) * incomingRadiance * dotNL;
}

vec4 texture2DArrayAA(texture2DArray tex, vec2 texCoord) {
    vec2 size = vec2(textureSize(sampler2DArray(tex, _Sampler), 0).xy);
    vec2 texCoordPX = texCoord * size;
    vec2 seam = floor(texCoordPX + vec2(0.5));

    texCoordPX = (texCoordPX - seam) / fwidth(texCoordPX) + seam;
    texCoordPX = clamp(texCoordPX, seam - 0.5, seam + 0.5);

    vec3 texCoord3 = vec3(texCoordPX / size, float(_TextureIndex));
    return texture(sampler2DArray(tex, _Sampler), texCoord3);
}

void main() {
    vec4 baseColorTexel = texture2DArrayAA(_BaseColorTexture, var_texCoord);
    vec4 occlusionRoughnessMetallicTexel = texture2DArrayAA(_OcclusionRoughnessMetallicTexture, var_texCoord);
    vec4 normalTextureTexel = texture2DArrayAA(_NormalTexture, var_texCoord);

    vec3 baseColor = baseColorTexel.rgb;
    float occlusion = occlusionRoughnessMetallicTexel.r;
    float roughness = occlusionRoughnessMetallicTexel.g;
    float metallic = occlusionRoughnessMetallicTexel.b;

    vec3 tangentVS = normalize(var_tangentVS);
    vec3 bitangentVS = normalize(var_bitangentVS);
    mat3 matrixTStoVS = mat3(tangentVS, bitangentVS, var_normalVS);
    vec3 normalTS = normalTextureTexel.xyz * 2.0 - 1.0;
    vec3 normalVS = normalize(matrixTStoVS * normalTS);

    vec3 positionVS = var_positionVS;
    vec3 viewDirectionVS = normalize(-positionVS);
    float dotNV = clamp(dot(normalVS, viewDirectionVS), 0.0, 1.0);
    float alpha = roughness * roughness;

    vec3 f0 = mix(DIELECTRIC_F0, baseColor, metallic);

    vec3 outgoingRadiance = vec3(0.0);

    for (int i = 0; i < _PointLightCount; i++) {
        Point_Light light = _PointLights[i];

        vec3 lightPositionVS = (_MatrixWStoVS * vec4(light.positionWS, 1.0)).xyz;
        vec3 lightDirectionVS = normalize(lightPositionVS - positionVS);
        float lightDistance = distance(positionVS, lightPositionVS);
        float lightAttenuation = 1.0 / (lightDistance * lightDistance);
        vec3 incomingRadiance = light.color * lightAttenuation;

        outgoingRadiance += lightOutgoingRadiance(
            viewDirectionVS, normalVS, dotNV,
            baseColor, alpha, metallic, f0,
            incomingRadiance, lightDirectionVS
        );
    }

    for (int i = 0; i < _DirectionalLightCount; i++) {
        Directional_Light light = _DirectionalLights[i];

        vec3 lightDirectionVS = normalize((_MatrixWStoVS * vec4(-light.directionWS, 0.0)).xyz);
        vec3 incomingRadiance = light.color;

        outgoingRadiance += lightOutgoingRadiance(
            viewDirectionVS, normalVS, dotNV,
            baseColor, alpha, metallic, f0,
            incomingRadiance, lightDirectionVS
        );
    }

    outgoingRadiance += _AmbientLight * baseColor * occlusion;

    vec3 toneMappedLinearColor = toneMapAcesNarkowicz(outgoingRadiance);
    vec3 toneMappedSrgbColor = pow(toneMappedLinearColor, vec3(1.0 / 2.2));

    fragColor = vec4(toneMappedSrgbColor, 1.0);
}

@end

@cs cs_equirectangular_to_cubemap

layout(binding = 0) uniform Layer_Index {
    int _LayerIndex;
};

layout(binding = 0) uniform texture2D _EquirectangularTexture;
layout(binding = 1, rgba16f) uniform writeonly image2DArray _CubemapImage;

layout(binding = 0) uniform sampler _EquirectangularSampler;

const float INV_PI = 0.31830987;
const float HALF_INV_PI = 0.15915494;

layout(local_size_x = 8, local_size_y = 8) in;
void main() {
    vec2 size = vec2(imageSize(_CubemapImage).xy);
    vec2 texCoord = (vec2(gl_GlobalInvocationID.xy) + vec2(0.5)) / size;

    texCoord = texCoord * vec2(2.0) - vec2(1.0); // Map to range [-1, 1]

    vec3 cubeCoord;
    if (_LayerIndex == 0) {
        // Positive X
        cubeCoord = vec3(1, -texCoord.y, -texCoord.x);
    } else if (_LayerIndex == 1) {
        // Negative X
        cubeCoord = vec3(-1, -texCoord.y, texCoord.x);
    } else if (_LayerIndex == 2) {
        // Positive Y
        cubeCoord = vec3(texCoord.x, 1, texCoord.y);
    } else if (_LayerIndex == 3) {
        // Negative Y
        cubeCoord = vec3(texCoord.x, -1, -texCoord.y);
    } else if (_LayerIndex == 4) {
        // Positive Z
        cubeCoord = vec3(texCoord.x, -texCoord.y, 1);
    } else if (_LayerIndex == 5) {
        // Negative Z
        cubeCoord = vec3(-texCoord.x, -texCoord.y, -1);
    }

    vec3 cubeDir = normalize(cubeCoord);

    float theta = atan(cubeDir.y, cubeDir.x);
    float phi = asin(cubeDir.z);

    vec2 equirectCoord = vec2(theta * HALF_INV_PI, phi * INV_PI) + vec2(0.5);

    vec4 irradiance = texture(sampler2D(_EquirectangularTexture, _EquirectangularSampler), equirectCoord);
    imageStore(_CubemapImage, ivec3(ivec2(gl_GlobalInvocationID.xy), _LayerIndex), irradiance);
}

@end

@program main vs_main fs_main
@program equirectangular_to_cubemap cs_equirectangular_to_cubemap