path tracer glsl shader

#version 430 core

#define EPS .000002
#define M_PI 3.1415926535897932384626433832795
#define SQRT_TWO 1.41421356237
#define spatialrand vec2
//#define CULL_BACKFACE

out vec4 FragColor;

uniform vec2 screenRes;

uniform mat4 camMat;
uniform vec3 camPos;

uniform sampler2D skyBox;

uniform int renderMode;

uniform int triTestThreshold;
uniform int boxTestThreshold;

uniform int frameCount;

uniform float time;

uniform int maxBounces;
uniform int raysPerPix;

uniform bool drawSkyBox;
uniform float skyBoxBrightness;
uniform bool dls;

uint triTests = 0;
uint boxTests = 0;

float weights[100];
float weightSum = 0;

bool singleRay = false;

vec3 aces(vec3 x)
{
  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((x * (a * x + b)) / (x * (c * x + d) + e), 0.0, 1.0);
}

vec3 acesInv(vec3 x) {
    return (sqrt(-10127.*x*x + 13702.*x + 9.) + 59.*x - 3.) / (502. - 486.*x);
}

float rand(inout uint state) {
    state = state * uint(747796405) + uint(2891336453);
    uint result = ((state >> ((state >> uint(28)) + uint(4))) ^ state) * uint(277803737);
    result = (result >> 22) ^ result;
    return result / 4294967295.0;
}

float randNormDist(inout uint state) {
    float theta = 2.0 * 3.1415926 * rand(state);
    float rho = sqrt(-2.0 * log(rand(state)));
    return abs(rho * cos(theta)) / 6;
}

vec3 isotropic(float rp, float c) {
  // sin(a) = sqrt(1.0 - cos(a)^2) , in the interval [0, PI/2] relevant for us
  float p = M_PI * 2 * rp, s = sqrt(1.0 - c*c);
  return vec3(cos(p) * s, sin(p) * s, c);
}

vec3 ortho(vec3 v) {
  return mix(vec3(-v.y, v.x, 0.0), vec3(0.0, -v.z, v.y), step(abs(v.x), abs(v.z)));
}

vec3 around(vec3 v, vec3 z) {
  vec3 t = ortho(z), b = cross(z, t);
  return fma(t, vec3(v.x), fma(b, vec3(v.y), z * v.z));
}

vec3 randomCosineWeightedHemispherePoint(vec3 n, inout uint state) {

  vec3 rand = vec3(rand(state), rand(state), rand(state));

  float c = sqrt(rand.y);
  return vec3(around(isotropic(rand.x, c), n));
}

float randRange(inout uint state, float minVal, float maxVal) {
    return rand(state) * (maxVal - minVal) + minVal;
}

struct Tri {
    vec4 v1;
    vec4 v2;
    vec4 v3;
    vec4 normal;
    vec4 minPos;
    vec4 maxPos;
    float area;
    float accumArea;
    float padding1;
    float padding2;
};

struct BoundingBox {
    vec4 minPos;
    vec4 maxPos;
    uint triIndex;
    uint triCount;
    uint childAIndex;
    uint childBIndex;
};

struct Ray {
    vec3 origin;
    vec3 dir;
    vec3 inLight;
    vec3 rayCol;
};

struct RayHit {
    bool hit;
    vec3 hitPos;
    vec3 normal;
    float dst;
    int modelIndex;
    int triIndex;
    bool internal;
};

struct Sphere {
    vec3 pos;
    float rad;
};

struct Material {
    vec4 diffuseCol;
    float roughness;
    float brightness;
    float padding1;
    float padding2;
};

struct Model {
    mat4 worldToModel;
    mat4 modelToWorld;
    Material mat;
    uint rootBoxIndex;
    float totalArea;
    float scale;
    uint lightIndex;
};

struct LightSelectResult {
    uint i;
    float weight;
    float weightSum;
};

layout (std430, binding = 0) buffer Triangles {
    Tri tris[];
};

layout (std430, binding = 1) buffer Boxes {
    BoundingBox boxes[];
};

layout (std430, binding = 2) buffer Models {
    Model models[];
};

layout (std430, binding = 3) buffer Colors {
    vec4 colors[];
};

layout (std430, binding = 4) buffer LightIdxs {
    unsigned int lightIdxs[];
};

/*layout (std430, binding = 2) buffer RootBoxIndexes {
    uint rootBoxIndex[];
};*/

Ray newRay(inout uint state, vec2 uv) {
    Ray r;
    //r.dir = normalize((vec4(uv, -1.0, 1.0) * camMat).xyz + vec3(randRange(state, 0.0, 0.1), randRange(state, 0.0, 0.1), 0.0));
    r.dir = normalize((vec4(uv, -1.0, 1.0) * camMat).xyz);
    r.dir += randomCosineWeightedHemispherePoint(r.dir, state) * 0.001;
    r.origin = camPos;
    r.rayCol = vec3(1.0);
    r.inLight = vec3(0.0);
    return r;
};

RayHit newRayHit() {
    RayHit r;
    r.hit = false;
    r.hitPos = vec3(0.0);
    r.normal = vec3(0.0);
    r.dst = 9999999.f;
    r.modelIndex = -1;
    r.triIndex = -1;
    r.internal = false;
    return r;
};

Sphere newSphere(vec3 pos, float rad)
{
    Sphere s;
    s.pos = pos;
    s.rad = rad;
    return s;
}

BoundingBox growToInclude(BoundingBox b, Tri t)
{
    b.minPos = vec4(min(b.minPos.x, t.v1.x), min(b.minPos.y, t.v1.y), min(b.minPos.z, t.v1.z), 0.0);
    b.maxPos = vec4(max(b.maxPos.x, t.v1.x), max(b.maxPos.y, t.v1.y), max(b.maxPos.z, t.v1.z), 0.0);

    b.minPos = vec4(min(b.minPos.x, t.v2.x), min(b.minPos.y, t.v2.y), min(b.minPos.z, t.v2.z), 0.0);
    b.maxPos = vec4(max(b.maxPos.x, t.v2.x), max(b.maxPos.y, t.v2.y), max(b.maxPos.z, t.v2.z), 0.0);

    b.minPos = vec4(min(b.minPos.x, t.v3.x), min(b.minPos.y, t.v3.y), min(b.minPos.z, t.v3.z), 0.0);
    b.maxPos = vec4(max(b.maxPos.x, t.v3.x), max(b.maxPos.y, t.v3.y), max(b.maxPos.z, t.v3.z), 0.0);

    return b;
}

float boxCollide(Ray r, BoundingBox b)
{
    vec3 invDir = 1.f / r.dir;
    vec3 tMin = (b.minPos.xyz - r.origin) * invDir;
    vec3 tMax = (b.maxPos.xyz - r.origin) * invDir;
    vec3 t1 = min(tMin, tMax);
    vec3 t2 = max(tMin, tMax);

    float dstFar = min(min(t2.x, t2.y), t2.z);
    float dstNear = max(max(t1.x, t1.y), t1.z);

    if (dstFar >= dstNear && dstFar > 0)
    {
        return dstNear;
    }
    ++boxTests;
    return 9999999.0;
}

RayHit triCollide(Ray ray, Tri tri)
{
    RayHit rHit = newRayHit();
    vec3 e1 = tri.v2.xyz - tri.v1.xyz; vec3 e2 = tri.v3.xyz - tri.v1.xyz;
    vec3 dxe2 = cross(ray.dir, e2);
    float det = dot(dxe2, e1);
    // ray parallel to triangle plane
    if (abs(det) < EPS) return rHit;
    if (det < EPS)
    {
        rHit.internal = true;
    }
    float invDet = 1. / det;
    vec3 t = ray.origin - tri.v1.xyz;
    float u = dot(dxe2, t) * invDet;
    vec3 txe1 = cross(t, e1);
    float v = dot(txe1, ray.dir) * invDet;
    if (u < 0. || v < 0. || u + v > 1.) return rHit;
    float d = dot(txe1, e2) * invDet;
    //result = vec3(d, u, v);
    if (d > 0.0)
    {
        rHit.hit = true;
        if (rHit.internal)
        {
            rHit.normal = -tri.normal.xyz;
        }
        else
        {
            rHit.normal = tri.normal.xyz;
        }
        rHit.dst = d;
        rHit.hitPos = ray.origin + ray.dir * d;
    }
    ++triTests;
    return rHit;
}

RayHit sphereCollide(Ray r, Sphere s) {
    RayHit rHit = newRayHit();

    float a = dot(r.dir, r.dir);
    vec3 s0_r0 = r.origin - s.pos;
    float b = 2.0 * dot(r.dir, s0_r0);
    float c = dot(s0_r0, s0_r0) - (s.rad * s.rad);
    float dst = (-b - sqrt((b*b) - 4.0*a*c))/(2.0*a);
    if (!(b*b - 4.0*a*c < 0.0) && dst > 0) {
        rHit.hit = true;
        rHit.dst = dst;
        rHit.hitPos = r.origin + r.dir * rHit.dst;
        rHit.normal = normalize(rHit.hitPos - s.pos);
    }
    return rHit;
}

RayHit findHit(Ray r) {

    RayHit closest = newRayHit();

    closest.hitPos = r.dir;

    uint stack[64];
    uint stackIndex = 0;

    Ray localRay = r;

    for (int modelIndex = 0; modelIndex < models.length(); ++modelIndex)
    {
        stack[stackIndex++] = models[modelIndex].rootBoxIndex;

        localRay.origin = (vec4(r.origin, 1.0) * models[modelIndex].worldToModel).xyz;
        localRay.dir = (vec4(r.dir, 0.0) * models[modelIndex].worldToModel).xyz;

        while (stackIndex > 0)
        {
            BoundingBox box = boxes[stack[--stackIndex]];

            if (box.childAIndex == 0)
            {
                for (uint i = box.triIndex; i < box.triIndex + box.triCount; ++i)
                {
                    RayHit rHit = triCollide(localRay, tris[i]);
                    //rHit.dst *= models[modelIndex].scale;

                    if (rHit.hit && rHit.dst < closest.dst)
                    {
                        closest = rHit;
                        closest.normal = normalize((vec4(closest.normal, 0.0) * models[modelIndex].modelToWorld).xyz);
                        closest.hitPos = (vec4(closest.hitPos, 1.0) * models[modelIndex].modelToWorld).xyz;
                        closest.modelIndex = modelIndex;
                        closest.triIndex = int(i);
                    }
                }
            }
            else
            {
                BoundingBox childA = boxes[box.childAIndex];
                BoundingBox childB = boxes[box.childBIndex];

                float dstA = boxCollide(localRay, childA);
                float dstB = boxCollide(localRay, childB);

                if (dstA > dstB)
                {
                    if (dstA < closest.dst) stack[stackIndex++] = box.childAIndex;
                    if (dstB < closest.dst) stack[stackIndex++] = box.childBIndex;
                }
                else
                {
                    if (dstB < closest.dst) stack[stackIndex++] = box.childBIndex;
                    if (dstA < closest.dst) stack[stackIndex++] = box.childAIndex;
                }
            }
        }
    }

    return closest;
}

vec2 sphericalMap(vec3 p)
{
    float theta = atan(p.x, p.z);

    vec3 vec = vec3(p.x, p.y, p.z);
    float radius = length(vec);
    float phi = acos(p.y / radius);

    float raw_u = theta / (2 * M_PI);

    float u = 1 - (raw_u + 0.5);

    float v = 1 - phi / (M_PI);

    return vec2(u, v);
}

/*void BRDF(inout Ray r, RayHit rH, inout uint state)
{
    if (models[rH.modelIndex].mat.brightness == 0)
    {
        r.dir = normalize((r.dir - 2.0 * (dot(r.dir, rH.normal)) * rH.normal) * (1 - models[rH.modelIndex].mat.roughness) + randomCosineWeightedHemispherePoint(rH.normal, state) * models[rH.modelIndex].mat.roughness);
        r.origin = rH.hitPos + r.dir * 0.001;
        r.rayCol = models[rH.modelIndex].mat.diffuseCol.xyz * r.rayCol;
    }
    else
    {

    }
}*/

uint sampleLightDistanceAreaWeighted(vec3 hitPos, inout uint state)
{
    //float weights[100];  // max 100 lights — adjust as needed
    weightSum = 0.0;

    for (uint i = 0; i < lightIdxs.length(); ++i) {
        vec3 lightPos = (vec4(0.0, 0.0, 0.0, 1.0) * models[lightIdxs[i]].modelToWorld).xyz;
        float distSq = dot(lightPos - hitPos, lightPos - hitPos) + EPS;
        weights[i] = models[lightIdxs[i]].mat.brightness * (models[lightIdxs[i]].totalArea * models[lightIdxs[i]].scale * models[lightIdxs[i]].scale) / distSq;
        weightSum += weights[i];
    }

    float cdf[100];
    float accum = 0.0;
    for (uint i = 0; i < lightIdxs.length(); ++i) {
        accum += weights[i] / weightSum;
        cdf[i] = accum;
    }

    float xi = rand(state);
    for (uint i = 0; i < lightIdxs.length(); ++i) {
        if (xi < cdf[i]) {
            return i;
        }
    }

    return uint(rand(state) * lightIdxs.length());
}

uint lightSampleTriIdx(uint modelIdx, inout uint state)
{
    float r = rand(state); // Uniform in [0, 1]

    // Binary search over CDF
    uint low = boxes[models[modelIdx].rootBoxIndex].triIndex;
    uint high = boxes[models[modelIdx].rootBoxIndex].triIndex + boxes[models[modelIdx].rootBoxIndex].triCount;
    while (low < high) {
        uint mid = (low + high) / 2;
        if (r < tris[mid].accumArea)
            high = mid;
        else
            low = mid + 1;
    }

    return low;
}

vec3 directLightSample(Ray r, RayHit rH, inout uint state)
{
    vec3 col = vec3(0.0);

    RayHit lH = newRayHit();

    uint i = sampleLightDistanceAreaWeighted(rH.hitPos, state);
    //uint i = uint(rand(state) * lightIdxs.length());;
    uint triIndex = lightSampleTriIdx(lightIdxs[i], state);

    float u = rand(state);
    float v = rand(state);

    if (u + v > 1.0f) {
        u = 1.0f - u;
        v = 1.0f - v;
    }

    vec4 v1 = tris[triIndex].v1 * models[lightIdxs[i]].modelToWorld;
    vec4 v2 = tris[triIndex].v2 * models[lightIdxs[i]].modelToWorld;
    vec4 v3 = tris[triIndex].v3 * models[lightIdxs[i]].modelToWorld;

    vec3 triNormal = (tris[triIndex].normal * models[lightIdxs[i]].modelToWorld).xyz;
    triNormal = normalize(triNormal);

    vec3 samplePoint = v1.xyz + u * (v2 - v1).xyz + v * (v3 - v1).xyz;

    vec3 toLight = samplePoint - rH.hitPos;
    float dstToLight = length(toLight);
    vec3 dir = normalize(toLight);

    float cos_surface = max(dot(rH.normal, dir), 0.0);
    float cos_light = max(dot(-dir, triNormal), 0.0);

    float lightPickPdf = weights[i] / weightSum;
    //float lightPickPdf = 1.0 / lightIdxs.length();

    float area =  models[lightIdxs[i]].totalArea / tris[triIndex].area;
    float pdfArea = (area * models[lightIdxs[i]].scale * models[lightIdxs[i]].scale);
    float pdfLight = ((dstToLight * dstToLight) / (cos_light * pdfArea)) * lightPickPdf;

    //pdfLight = 1.0;

    float pdfBRDF = max(dot(dir, rH.normal), 0.0) / M_PI;

    Ray r2 = r;
    r2.dir = dir;

    r2.origin = rH.hitPos + r2.dir * 0.001;

    lH = findHit(r2);
    if (lH.hit && lH.triIndex == triIndex && lH.modelIndex == lightIdxs[i])
    {
        vec3 emittedLight = models[lightIdxs[i]].mat.diffuseCol.xyz * models[lightIdxs[i]].mat.brightness;// Emission per triangle
        vec3 brdf = models[rH.modelIndex].mat.diffuseCol.xyz / M_PI;

        vec3 contribution = (emittedLight * brdf * cos_surface) / (pdfLight);
        //vec3 contribution = (emittedLight * brdf * cos_surface * cos_light) / (dstToLight * dstToLight * pdfLight);

        //float w = (pdfLight * pdfLight) / (pdfLight * pdfLight + pdfBRDF * pdfBRDF);
        float w = pdfLight / (pdfLight + pdfBRDF);

        col += contribution * r.rayCol * w;
    }
    return col;
}

vec4 trace(Ray r, inout uint state)
{
    RayHit rH = newRayHit();

    vec3 emittedLight = vec3(0.0);

    for (int bounces = 0; bounces < maxBounces; ++bounces)
    {
        rH = findHit(r);

        if (rH.hit)
        {
            if (models[rH.modelIndex].mat.brightness == 0)
            {
                if (dls)
                {
                    r.inLight += directLightSample(r, rH, state);
                    //return vec4(directLightSample(r, rH, state), 1.0);
                }

                r.dir = normalize((r.dir - 2.0 * (dot(r.dir, rH.normal)) * rH.normal) * (1 - models[rH.modelIndex].mat.roughness) + randomCosineWeightedHemispherePoint(rH.normal, state) * models[rH.modelIndex].mat.roughness);
                r.origin = rH.hitPos + r.dir * 0.001;
                r.rayCol = models[rH.modelIndex].mat.diffuseCol.xyz / M_PI * r.rayCol;

            }
            else
            {
                emittedLight = models[rH.modelIndex].mat.diffuseCol.xyz * models[rH.modelIndex].mat.brightness;
                if (bounces == 0)
                {
                    singleRay = true;
                    return vec4(emittedLight, 1.0);
                    //return vec4(vec3(tris[rH.triIndex].accumArea), 1.0);
                }

                if (dls)
                {
                    float pdfBRDF = max(dot(r.dir, rH.normal), 0.0);

                    float cos_surface = max(dot(rH.normal, r.dir), 0.0);
                    float cos_light = max(dot(-r.dir, tris[rH.triIndex].normal.xyz), 0.01); // clamp here!

                    uint light = models[rH.modelIndex].lightIndex;

                    float lightPickPdf = weights[light] / weightSum;
                    //float pdfArea = 1.0 / (models[rH.modelIndex].totalArea * models[rH.modelIndex].scale * models[rH.modelIndex].scale);
                    //float pdfLight = ((rH.dst * rH.dst) / cos_light) * pdfArea * lightPickPdf;
                    float area =  models[rH.modelIndex].totalArea / tris[rH.triIndex].area;
                    float pdfArea = (area * models[rH.modelIndex].scale * models[rH.modelIndex].scale);
                    float pdfLight = ((rH.dst * rH.dst) / (cos_light * pdfArea)) * lightPickPdf;
                    //float pdfLight = pdfArea * rH.dst * rH.dst / (cos_light + 0.000001);

                    //float w = (pdfBRDF * pdfBRDF) / (pdfLight * pdfLight + pdfBRDF * pdfBRDF);
                    float w = pdfBRDF / (pdfLight + pdfBRDF);

                    r.inLight += w * emittedLight * r.rayCol;
                }
                else
                {
                    r.inLight += emittedLight * r.rayCol;
                }

                return vec4(r.inLight, 1.0);
            }
        }
        else
        {
            if (bounces == 0)
            {
                singleRay = true;
                if (drawSkyBox)
                {
                    r.inLight += r.rayCol * acesInv(texture(skyBox, sphericalMap(-r.dir)).xyz * skyBoxBrightness);
                    return vec4(r.inLight, 1.0);
                }
                return vec4(0.0);
            }
            else
            {
                if (drawSkyBox)
                {
                    r.inLight += r.rayCol * acesInv(texture(skyBox, sphericalMap(-r.dir)).xyz * skyBoxBrightness);
                    return vec4(r.inLight, 1.0);
                }
                return vec4(r.inLight, 1.0);
            }
        }
    }
    return vec4(r.inLight, 1.0);
}

void main()
{
    vec2 texUV = gl_FragCoord.xy / screenRes;
    vec2 uv = ((texUV - 0.5) * 2) * vec2(screenRes.x / screenRes.y, 1.0);
    uint pixelIndex = uint(gl_FragCoord.y * screenRes.x + gl_FragCoord.x);
    uint rngState = pixelIndex + uint(time * 3000 * (2508 - time));

    Ray r = newRay(rngState, uv);

    switch (renderMode) {
        case 0:
            RayHit rHit = findHit(r);
            //FragColor = vec4(0.0, 1.0, 0.0, 1.0);

            if (rHit.hit) {
                FragColor = models[rHit.modelIndex].mat.diffuseCol * (dot(rHit.normal, vec3(1.0)) + 1) * 0.5;
            }
            else {
                if (drawSkyBox) {
                    FragColor = texture(skyBox, sphericalMap(-r.dir));
                }
                else {
                    FragColor = vec4(0.0, 0.0, 0.0, 1.0);
                }
            }

            break;
        case 1:

            if (frameCount == 1)
            {
                vec4 col = vec4(0.0);
                for (int i = 0; i < raysPerPix; ++i)
                {
                    col += trace(r, rngState);
                    if (singleRay)
                    {
                        col *= raysPerPix;
                        break;
                    }
                }
                colors[pixelIndex] = col;
                FragColor = vec4(aces(colors[pixelIndex].xyz / raysPerPix), 1.0);
                //FragColor = vec4(colors[pixelIndex].xyz, 1.0);
                //FragColor = vec4(vec3(a) / raysPerPix, 1.0);
            }
            else
            {
                vec4 col = vec4(0.0);
                for (int i = 0; i < raysPerPix; ++i)
                {
                    col += trace(r, rngState);
                    if (singleRay)
                    {
                        col *= raysPerPix;
                        break;
                    }
                }
                colors[pixelIndex] += col;
                FragColor = vec4(aces(colors[pixelIndex].xyz / (raysPerPix * frameCount)), 1.0);
                //FragColor = vec4(colors[pixelIndex].xyz / frameCount, 1.0);
                //FragColor = vec4(vec3(a) / raysPerPix, 1.0);
            }
            break;

        case 2:
            vec4 col1 = vec4(0.0);
            for (int i = 0; i < raysPerPix; ++i)
            {
                col1 += trace(r, rngState);
                if (singleRay)
                {
                    break;
                }
            }
            if (triTests > triTestThreshold)
            {
                FragColor = vec4(1.0, 0.0, 0.0, 1.0);
            }
            else
            {
                FragColor = vec4(vec3(float(triTests) / float(triTestThreshold)), 1.0);
            }
            break;

        case 3:
            vec4 col2 = vec4(0.0);
            for (int i = 0; i < raysPerPix; ++i)
            {
                col2 += trace(r, rngState);
                if (singleRay)
                {
                    break;
                }
            }
            if (boxTests > boxTestThreshold)
            {
                FragColor = vec4(1.0, 0.0, 0.0, 1.0);
            }
            else
            {
                FragColor = vec4(vec3(float(boxTests) / float(boxTestThreshold)), 1.0);
            }
            break;
    }
}