precision highp float;

uniform highp vec2 iRotate;
uniform highp float iDist;
uniform highp vec2 iResolution;

const vec3 light = normalize(vec3(0, 0, 1));

bool intersect_sphere(vec3 ce, float r, in vec3 ro, in vec3 rd, out float t, out vec3 n) {
    vec3 p = ro-ce;
    float b = dot(p,rd), c = dot(p,p)-r*r;
    float delta = b*b-c; if (delta<=0.0) return false;
    delta = sqrt(delta);
    t = -b-delta; if (t<=0.0) t = -b+delta;
    if (t<=0.0) return false;
    n = (p+rd*t)/r; return true;
}

bool intersect_rod(vec3 pa, vec3 pb, float r, in vec3 ro, in vec3 rd, out float t, out vec3 n) {
    vec3 ab = pb-pa, p = ro-pa;
    float ab2 = dot(ab, ab), abrd = dot(ab, rd), abp = dot(ab, p);
    float a = ab2 - abrd*abrd;
    float b = ab2*dot(p,rd) - abp*abrd;
    float c = ab2*dot(p,p) - abp*abp - r*r*ab2;
    float delta = b*b-a*c; if (delta<=0.0) return false;
    delta = sqrt(delta);
    t = (-b-delta)/a;
    if (t>0.0) {
        float h = abp+t*abrd;
        if (h>0.0 && h< ab2) {
            n = (p+rd*t-ab*h/ab2)/r;
            return true;
        }
    }
    t = (-b+delta)/a;
    if (t>0.0) {
        float h = abp+t*abrd;
        if (h>0.0 && h< ab2) {
            n = (p+rd*t-ab*h/ab2)/r;
            return true;
        }
    }
    return false;
}

vec3 traceRay(vec3 ro, vec3 rd) {
    
    const float R = 0.2, r = 0.1;
    const float rt3_2 = 0.8660254037844386;

    vec3 t_col = vec3(1.0), f_col;

    for (int i = 0; i < 8; i++) {
        ro += 1e-3*rd;
        float min_t = 1e+12, t;
        vec3 min_n = vec3(0.0), n;

        // intersect with the plane
        t = -(ro.z+(rt3_2+R))/rd.z;
        if (t > 0.0) {
            min_t = t;
            f_col = vec3(0.5, 0.8, 1.0);
            vec2 p = ro.xy+rd.xy*t;
            if (mod(floor(p.x)+floor(p.y),2.0)==0.0) f_col*=0.99;
            min_n = vec3(0, 0, 1);
        }

        // intersect with balls
        if (intersect_sphere(vec3(0, 0, 0), R, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = vec3(1.0);
        }
        if (intersect_sphere(vec3(1, 0, 0), R, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = vec3(1.0);
        }
        if (intersect_sphere(vec3(0, 1, 0), R, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = vec3(1.0);
        }
        if (intersect_sphere(vec3(0.5, -0.5, rt3_2), R, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = vec3(1.0);
        }
        if (intersect_sphere(vec3(0.5, 0.5, rt3_2), R, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = vec3(1.0);
        }
        if (intersect_sphere(vec3(-0.5, 0.5, rt3_2), R, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = vec3(1.0);
        }
        if (intersect_sphere(vec3(0.5, 0.5, -rt3_2), R, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = vec3(1.0);
        }

        // intersect with rods
        const vec3 rod_col = vec3(1.0, 0.7, 0.75);
        if (intersect_rod(vec3(-1, 0, 0), vec3(1, 0, 0), r, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = rod_col;
        }
        if (intersect_rod(vec3(0, -1, 0), vec3(0, 1, 0), r, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = rod_col;
        }
        if (intersect_rod(vec3(-0.5, -0.5, -rt3_2), vec3(0.5, 0.5, rt3_2), r, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = rod_col;
        }
        if (intersect_rod(vec3(-0.5, 0.5, -rt3_2), vec3(0.5, -0.5, rt3_2), r, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = rod_col;
        }
        if (intersect_rod(vec3(0.5, -0.5, -rt3_2), vec3(-0.5, 0.5, rt3_2), r, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = rod_col;
        }
        if (intersect_rod(vec3(-0.5, -0.5, rt3_2), vec3(0.5, 0.5, -rt3_2), r, ro, rd, t, n) && t < min_t) {
            min_t = t, min_n = n, f_col = rod_col;
        }

        if (min_n==vec3(0.0)) break;

        rd = normalize(rd), min_n = normalize(min_n);  // prevent error accumulation

        ro = ro + rd*min_t;
        rd = rd - 2.0*dot(rd, min_n)*min_n;
        t_col *= f_col;
    }

    t_col *= max(dot(rd, light), 0.0);
    return t_col;
}

#define MAX_AA 5
uniform int AA;

void main() {

    float rx = iRotate.x, rz = iRotate.y;
    vec3 w = vec3(cos(rx)*vec2(cos(rz),sin(rz)), sin(rx));
    vec3 u=vec3(-sin(rz),cos(rz),0);
    vec3 v=cross(w,u);

    vec3 ro = iDist*w;

    vec3 col = vec3(0.0);
    for (int i=0;i< MAX_AA;i++) {
        for (int j=0;j< MAX_AA;j++) {
            vec2 aVertexPosition = 2.0 * (gl_FragCoord.xy + vec2(i,j)/float(AA)) / iResolution - vec2(1.0);
            vec3 rd = normalize(mat3(u,v,-w)*vec3(aVertexPosition*iResolution.xy,2.0*length(iResolution)));
            col += clamp(traceRay(ro, normalize(rd)), 0.0, 1.0);
            if (j==AA-1) break;
        }
        if (i==AA-1) break;
    }
    col /= float(AA*AA);

    gl_FragColor = vec4(col,1.0);
}