const int DEPTH = 32;
const int AA = 8;

const vec2 epsilon = vec2(1e-3, 0);

const float planeNear = 1e-3;
const float planeFar = 1e2;
const float fov = 82.0;

const float PI = 3.141592653589793;

const float pdf = 2.0 * PI;

#define UNION(a, b) ((a.x) < (b.x) ? (a) : (b))
#define SDF vec2

#define MATERIAL_WHITE 0
#define MATERIAL_RED 1
#define MATERIAL_GREEN 2
#define MATERIAL_BLUE 3
#define MATERIAL_MIRROR 10
#define MATERIAL_GLASS 20
#define MATERIAL_EMISSION 255

SDF sdfScene(in vec3 pos)
{
    vec3 q = abs(pos - vec3(0,0.99,-1)) - vec3(0.5, 0.01, 0.5);
    SDF li0 = SDF(length(max(q, 0.0)) + min(max(q.x, max(q.y, q.z)), 0.0), MATERIAL_EMISSION);

    SDF sp1 = SDF(length(pos + vec3( 0.5, 0.5, 1)) - 0.5, MATERIAL_WHITE);
    SDF sp2 = SDF(abs(length(pos + vec3(-0.5, 0.75, 1)) - 0.25) - epsilon[0], MATERIAL_GLASS);
    SDF sp3 = SDF(length(pos + vec3( 0.5,-0.25, 1)) - 0.25, MATERIAL_MIRROR);

    SDF pl0 = SDF( pos.y + 1.0, MATERIAL_WHITE);
    SDF pl1 = SDF(-pos.y + 1.0, MATERIAL_WHITE);
    SDF pl2 = SDF( pos.x + 1.0, MATERIAL_RED);
    SDF pl3 = SDF(-pos.x + 1.0, MATERIAL_GREEN);
    SDF pl4 = SDF(-pos.z,       MATERIAL_BLUE);
    SDF pl5 = SDF( pos.z + 4.1, MATERIAL_WHITE);
    
    SDF box = UNION(pl0, pl1);
        box = UNION(box, pl2);
        box = UNION(box, pl3);
        box = UNION(box, pl4);
        box = UNION(box, pl5);

    box = UNION(box, sp1);
    box = UNION(box, sp2);
    box = UNION(box, sp3);

    return UNION(li0, box);
}

vec3 sceneNormal(in vec3 pos)
{
    return normalize(sdfScene(pos)[0] - vec3(
        sdfScene(pos + epsilon.xyy)[0],
        sdfScene(pos + epsilon.yxy)[0],
        sdfScene(pos + epsilon.yyx)[0]
    ));
}

// project vector b onto vector a
vec3 project(in vec3 a, in vec3 b)
{
    return a * (dot(a, b) / dot(a, a));
}

// construct a 3D coordinate system with the input up being the "upwards" facing vector
// which will be directly stored in w.
// Mathematically this function will create two non linear vectors of up and generate an orthonormal
// basis using gram-schmidt.
// This function assumes "up" being already normalized
void construct_orthonormal_basis(in vec3 up, out vec3 u, out vec3 v, out vec3 w)
{
    w = up;

    vec3 n2 = normalize(cross(w, vec3(0, 1.0, 1.0))); // build perpendicular vector from w
    vec3 n3 = cross(w, n2); // create 2nd vector perpendicular to w and n2

    // gram schmidt
    u = n2 - (project(w, n2));
    v = n3 - project(w, n3) - project(u, n3);
}

// generate a normalized vector within the bounds of the hemisphere with radius of 1.
// Z-Coordinate will be "upwards".
// radius determines the maximum radius the output vector will have.
// NOTE: shrinkin radius will still result in the output to be normalized
vec3 cosine_weighted_hemisphere()
{
    float u1 = random();
    float u2 = random();

    float r = sqrt(u1);
    float theta = 2.0 * PI * u2;

    float x = r * cos(theta);
    float y = r * sin(theta);

    return vec3(x, y, sqrt(1.0 - u1));
}

vec3 generate_brdf_ray_direction(in vec3 normal, in vec3 incident) {
    vec3 hemisphere = cosine_weighted_hemisphere();

    vec3 u, v, w;
    construct_orthonormal_basis(normal, u, v, w);

    return u * hemisphere.x + v * hemisphere.y + w * hemisphere.z;
}

vec3 tracePath(in vec3 ro, in vec3 rd)
{
    vec3 col = vec3(0);
    vec3 thr = vec3(1);

    for (int k = 0; k < DEPTH; k++)
    {
        vec3 pos = ro + rd * planeNear;
        float t = planeNear;
        for (int i = 0; i < 256 && t < planeFar; i++)
        {
            SDF h = sdfScene(pos);

            if (h[0] < epsilon[0])
            {
                vec3 nor = sceneNormal(pos);

                switch(int(h[1]))
                {
                    case MATERIAL_MIRROR:
                        rd = reflect(rd, nor);
                        ro = pos - nor * epsilon[0];
                        break;
                    case MATERIAL_GLASS:
                        if (dot(nor, rd) > random())
                        {
                            rd = refract(rd, -nor, 0.2);
                            ro = pos + nor * epsilon[0] * 4.0;
                        }
                        else
                        {
                            rd = reflect(rd, nor);
                            ro = pos - nor * epsilon[0];
                        }
                        break;
                    default:
                        rd = -generate_brdf_ray_direction(nor, rd);
                        ro = pos - nor * epsilon[0];
                        break;
                }

                float cosTheta = abs(dot(nor, -rd)) * pdf / PI * 0.6667;

                switch(int(h[1]))
                {
                    case MATERIAL_EMISSION:
                        col += thr * 10.0;
                        break;
                    case MATERIAL_WHITE:
                        thr *= cosTheta;
                        break;
                    case MATERIAL_RED:
                        thr *= vec3(0.9, 0.1, 0.1) * cosTheta;
                        break;
                    case MATERIAL_GREEN:
                        thr *= vec3(0.1, 0.9, 0.1) * cosTheta;
                        break;
                    case MATERIAL_BLUE:
                        thr *= vec3(0.1, 0.1, 0.9) * cosTheta;
                        break;
                    default:
                        break;
                }
                break;
            }

            t += h[0];
            pos += rd * h[0];
        }

        float m = max(thr.x, max(thr.y, thr.z));
        if (random() > m)
            break;
    }
    return col;
}

void mainImage(out vec4 fragColor, in vec2 fragCoord)
{
    SEED = (fragCoord.x * fragCoord.y)/iResolution.y + iTime;

    vec3 col = vec3(0);
    vec3 ro = vec3(0,0,-3.999);

    // take multiple samples per pixel
    // offset each ray direction slightly inside pixel
    // in order to anti aliase
    for (int i = 0; i < AA; i++, SEED++)
    {
        vec2 off = vec2(random(), random()) * 0.5;
        vec2 uv = (fragCoord + off -0.5 * iResolution.xy) / iResolution.y;
        vec3 rd = normalize(vec3(uv, 1.0/tan(fov * 0.008726646259971648)));

        col += tracePath(ro, rd);
    }
    col /= float(AA);
    
    vec3 prev = texture(iChannel0, fragCoord.xy/iResolution.xy).rgb;
    col = mix(prev, col, 1.0/max(float(iFrame), 1.0));

    fragColor = vec4(col, 1.0);
}