raytracer.comp

#version 430 core

#define M_PI 3.1415926535897932384626433832795

#define NUM_SPHERE 4
#define NUM_MATERIAL 9

#define NX 800
#define NY 600

#define NS_X 2
#define NS_Y 1

#define MAX_DEPTH 10

layout (local_size_x = 2, local_size_y = 3) in;

// layout (rgba32f, binding = 0) uniform image2D specular0;
// layout (rgba32f, binding = 2) uniform image2D diffuse0;
// layout (rgba32f, binding = 4) uniform image2D ambient0;

// layout (rgba32f, binding = 1) uniform image2D specular1;
// layout (rgba32f, binding = 3) uniform image2D diffuse1;
// layout (rgba32f, binding = 5) uniform image2D ambient1;

layout (rgba32f, binding = 0) uniform image2D specular0;
layout (rgba32f, binding = 2) uniform image2D diffuse0;
layout (rgba32f, binding = 4) uniform image2D ambient0;

layout (rgba32f, binding = 1) uniform image2D specular1;
layout (rgba32f, binding = 3) uniform image2D diffuse1;
layout (rgba32f, binding = 5) uniform image2D ambient1;

// uniform sampler2D tex0;
layout (binding = 6) uniform sampler2D tex0;
layout (binding = 7) uniform sampler2D tex1;
layout (binding = 8) uniform sampler2D tex2;
layout (binding = 9) uniform sampler2D tex3;
layout (binding = 10) uniform sampler2D tex4;
layout (binding = 11) uniform sampler2D tex5;
layout (binding = 12) uniform sampler2D tex6;
layout (binding = 13) uniform sampler2D tex7;
layout (binding = 14) uniform sampler2D tex8;

uniform vec3 cam_origin;
uniform vec3 cam_lookup;
uniform vec3 cam_lookat;


// orbit r: vec4.x
// self r: vec4.y
// orbit degree: vec4.z
// rotate degree: vec4.w
uniform vec4 sphere0;
uniform vec4 sphere1;
uniform vec4 sphere2;
uniform vec4 sphere3;
uniform vec4 sphere4;
uniform vec4 sphere5;
uniform vec4 sphere6;
uniform vec4 sphere7;
uniform vec4 sphere8;

const vec3 ambient_light = vec3(1.0, 1.0, 1.0);

struct Ray{
    vec3 from;
    vec3 dir;
};

struct Rec4Hit{
    bool if_hit;
    float t;
    vec3 norm;
    vec3 p;
    vec2 tex_xy;
    int mat_idx;
};

struct Rec4Reflect{
    bool if_reflect;
    Ray reflected;
    vec3 attenuation;
};

struct Sphere{
    vec3 c;
    float r;
    float deg;
    int mat_idx;
} list_sphere[NUM_SPHERE];

struct Material{
    vec3 color;
    int tex_idx;
    float emit;
    float specular;
    float diffuse;
    float ambient;
} list_material[NUM_MATERIAL];

struct Camera{
    vec3 horizontal;    // FOV horizontal axis(X)
    vec3 vertical;  // FOV vertical axis(Y)
    vec3 origin;    // Camera location
    vec3 center;    // FOV center
} cam;

vec2 random_seed;

void setSeed(vec2 seed){
    random_seed = seed;
}

float _random(vec2 co){
    return fract(sin(dot(co.xy ,vec2(12.9898,78.233))) * 43758.5453);
}

float random01(){
    //return a random float number in [0, 1)

    random_seed.x = _random(random_seed);
    random_seed.y = _random(random_seed);

    return random_seed.x;
} 

vec3 randomUnitV3(){
    float x = cos(radians(180 * random01()));
    float y = cos(radians(180 * random01()));
    float z = sqrt(1 - x * x - y * y);

    return vec3(x, y, z);
}

vec3 pointAt(Ray ray, float t){
    return ray.from + t * ray.dir;
}

vec2 getSphereTexXY(vec3 p, Sphere sp, float deg){
    p = (p - sp.c) / sp.r;

    float cosine = cos(radians(deg));
    float sine = sin(radians(deg));

    float radian_x = atan(dot(p, vec3(cosine, 0.0, -sine)),
                          dot(p, vec3(sine, 0.0, cosine)));
    float x = radian_x / M_PI + 0.5;

    float y = p.y + 0.5;

    return vec2(x, y);
}

Rec4Hit hitSphere(Ray ray, Sphere sp, float t_min, float t_max){
    float a = dot(ray.dir, ray.dir);
    float b = 2 * dot(ray.dir, ray.from - sp.c);
    float c = dot(ray.from - sp.c, ray.from - sp.c) - sp.r * sp.r;

    float delta = b * b - 4 * a * c;

    Rec4Hit rec;
    rec.mat_idx = sp.mat_idx;
    if (delta >= 0){
        rec.if_hit = true;
        
        float tmp = (- b - sqrt(delta)) / (2 * a);
        
        if (t_min <= tmp && tmp <= t_max){
            rec.t = tmp;
            rec.p = pointAt(ray, tmp);

            rec.norm = normalize(rec.p - sp.c);

            rec.p = sp.c + rec.norm * sp.r;
            rec.tex_xy = getSphereTexXY(rec.p, sp, sp.deg);

            return rec;
        }

        tmp = (- b + sqrt(delta)) / (2 * a);

        if (t_min <= tmp && tmp <= t_max){
            rec.t = tmp;
            rec.p = pointAt(ray, tmp);

            rec.norm = normalize(rec.p - sp.c);

            rec.p = sp.c + rec.norm * sp.r;
            rec.tex_xy = getSphereTexXY(rec.p, sp, sp.deg);

            return rec;
        }
    }

    rec.if_hit = false;
    return rec;
}

Rec4Hit hitSphereList(Ray ray, Sphere list[NUM_SPHERE], float t_min, float t_max){
    Rec4Hit rec;
    rec.if_hit = false;

    for (int i = 0; i < NUM_SPHERE; i++){
        Rec4Hit tmp = hitSphere(ray, list[i], t_min, t_max);

        if (tmp.if_hit){
            t_max = tmp.t;
            rec = tmp;
        }
    }

    return rec;
}

vec3 getSurfaceColor(Material mat, vec2 tex_xy){
    switch (mat.tex_idx){
        case 0:
            return texture(tex0, tex_xy).xyz;
            break;
        case 1:
            return texture(tex1, tex_xy).xyz;
            break;
        case 2:
            return texture(tex2, tex_xy).xyz;
            break;
        case 3:
            return texture(tex3, tex_xy).xyz;
            break;
        case 4:
            return texture(tex4, tex_xy).xyz;
            break;
        case 5:
            return texture(tex5, tex_xy).xyz;
            break;
        case 6:
            return texture(tex6, tex_xy).xyz;
            break;
        case 7:
            return texture(tex7, tex_xy).xyz;
            break;
        case 8:
            return texture(tex8, tex_xy).xyz;
            break;
    }

    return mat.color;
}

Rec4Reflect specular(Rec4Hit rec, Ray ray){
    Rec4Reflect rec_rf;

    if (list_material[rec.mat_idx].emit > 0){
        rec_rf.if_reflect = false;
        return rec_rf;
    }

    vec3 dir_in = normalize(ray.dir);
    vec3 new_dir = ray.dir -  2 * dot(dir_in, rec.norm) * rec.norm;
    rec_rf.reflected = Ray(rec.p, new_dir);

    rec_rf.attenuation = getSurfaceColor(list_material[rec.mat_idx], rec.tex_xy) * list_material[rec.mat_idx].specular;

    rec_rf.if_reflect = dot(new_dir, rec.norm) > 0;

    return rec_rf;
}

Rec4Reflect diffuse(Rec4Hit rec, Ray ray){
    Rec4Reflect rec_rf;

    if (list_material[rec.mat_idx].emit > 0){
        rec_rf.if_reflect = false;
        return rec_rf;
    }

    vec3 new_dir = randomUnitV3() * random01() + rec.norm;
    rec_rf.reflected = Ray(rec.p, new_dir);

    rec_rf.attenuation = getSurfaceColor(list_material[rec.mat_idx], rec.tex_xy) * list_material[rec.mat_idx].diffuse;

    // rec_rf.if_reflect = dot(new_dir, rec.norm) > 0;
    rec_rf.if_reflect = true;

    return rec_rf;
}

Rec4Reflect ambient(Rec4Hit rec, Ray ray){
    Rec4Reflect rec_rf;

    if (list_material[rec.mat_idx].emit > 0){
        rec_rf.if_reflect = false;
        return rec_rf;
    }

    rec_rf.attenuation = getSurfaceColor(list_material[rec.mat_idx], rec.tex_xy) * list_material[rec.mat_idx].ambient;
    rec_rf.if_reflect = true;

    return rec_rf;
}

vec3 emit(Rec4Hit rec, Ray ray){
    // float decay = sp.r / distance(ray.from, rec.p);
    // decay *= decay;

    float cosine = -dot(rec.norm, normalize(ray.dir));

    return getSurfaceColor(list_material[rec.mat_idx], rec.tex_xy) * list_material[rec.mat_idx].emit;
}

vec3 color(Ray r, int reflect_type){
    vec3 total_att = vec3(1.0, 1.0, 1.0);
    vec3 final_color = vec3(0.0, 0.0, 0.0);

    Ray ray = r;

    for (int i = 0; i < MAX_DEPTH; i++){
        Rec4Hit rec = hitSphereList(ray, list_sphere, 0.0001, 10000000);

        if (!rec.if_hit){
            break;
        }
        else{
            Rec4Reflect rec_rf;

            switch (reflect_type){
                case 0:
                    rec_rf = specular(rec, ray);
                    break;
                case 1:
                    rec_rf = diffuse(rec, ray);
                    break;
                case 2:
                    rec_rf = ambient(rec, ray);
                    break;
            }

            if (!rec_rf.if_reflect){
                final_color = emit(rec, ray);
                break;
            }

            if (reflect_type == 2){
                return ambient_light * rec_rf.attenuation;
            }

            ray = rec_rf.reflected;

            total_att *= rec_rf.attenuation;
        }
    }

    return total_att * final_color;
}

void saveTex(vec3 texel, ivec2 coord, int type){
    switch (type){
        case 0:
            imageStore(specular0, coord, vec4(texel, 1.0));
            break;
        case 1:
            imageStore(specular1, coord, vec4(texel, 1.0));
            break;

        case 2:
            imageStore(diffuse0, coord, vec4(texel, 1.0));
            break;
        case 3:
            imageStore(diffuse1, coord, vec4(texel, 1.0));
            break;

        case 4:
            imageStore(ambient0, coord, vec4(texel, 1.0));
            break;
        case 5:
            imageStore(ambient1, coord, vec4(texel, 1.0));
            break;
    }
}

vec2 rotateVec2(vec2 v, float rad){
    float l = length(v);
    float r = acos(normalize(v).x);
    return vec2(cos(r + rad), sin(r + rad) * l);
}

Sphere genSphere(vec4 v){
    Sphere sp;
    sp.c = vec3(cos(radians(v.z)), 0.0, - sin(radians(v.z))) * v.x;
    sp.r = v.y;
    sp.deg = v.w;

    return sp;
}

void initSpheres(){
    list_sphere[0] = genSphere(sphere0);
    list_sphere[1] = genSphere(sphere1);
    list_sphere[2] = genSphere(sphere2);
    list_sphere[3] = genSphere(sphere3);

    for (int i = 0; i < NUM_SPHERE; i++){
        list_sphere[i].mat_idx = i;
    }
}


//emit: v.x
//specular: v.y
//diffuse: v.z
//ambient: v.y + v.z;
Material genMaterial(vec3 v){
    Material mat;
    mat.emit = v.x;

    mat.specular = v.y;
    mat.diffuse = v.z;
    mat.ambient = v.y + v.z;

    return mat;
}

void initMaterials(){

    list_material[0] = genMaterial(vec3(1.0, 0.0, 0.0));
    list_material[1] = genMaterial(vec3(0.0, 0.6, 0.35));
    list_material[2] = genMaterial(vec3(0.0, 0.6, 0.35));
    list_material[3] = genMaterial(vec3(0.0, 0.6, 0.35));
    list_material[4] = genMaterial(vec3(0.0, 0.35, 0.6));
    list_material[5] = genMaterial(vec3(0.0, 0.35, 0.6));
    list_material[6] = genMaterial(vec3(0.0, 0.35, 0.6));
    list_material[7] = genMaterial(vec3(0.0, 0.35, 0.6));
    list_material[8] = genMaterial(vec3(0.0, 0.35, 0.6));

    for (int i = 0; i < NUM_MATERIAL; i++){
        list_material[i].tex_idx = i;
    }
}

vec2 convertCoord(vec2 tex_coord, int width, int height){
    vec2 coord = tex_coord;
    
    coord -= vec2(width, height) / 2;
    coord /= width;

    return coord;
}

void initCamera(vec3 origin, vec3 lookat, vec3 lookup){
    cam.horizontal = normalize(cross(lookat, lookup));
    cam.vertical = normalize(lookup);
    cam.origin = origin;
    cam.center = origin + normalize(lookat);
}

Ray getRay(Camera cam, vec2 coord){
    vec3 shootat = cam.center;

    shootat += cam.horizontal * coord.x;
    shootat += cam.vertical * coord.y;

    return Ray(cam.origin, shootat - cam.origin);
}

vec3 fuck(Ray ray, int reflect_type){
    switch(reflect_type){
        case 0:
            return vec3(1.0, 0.0, 0.0);
        case 1:
            return vec3(0.0, 1.0, 0.0);
        case 2:
            return vec3(0.0, 0.0, 1.0);
        case 3:
            return vec3(1.0, 1.0, 0.0);
        case 4:
            return vec3(0.0, 1.0, 1.0);
        case 5:
            return vec3(1.0, 0.0, 1.0);
        
    }
    return vec3(0.5, 0.5, 0.5);
}

void main(){
    int reflect_type = int(gl_LocalInvocationID.y);
    int map_id = 2 * int(gl_LocalInvocationID.y) + int(gl_LocalInvocationID.x);
    ivec2 texCoord = ivec2(gl_WorkGroupID.xy);

    setSeed(convertCoord(texCoord, NX, NY));

    initSpheres();
    initMaterials();
    initCamera(cam_origin,  //origin
               cam_lookat,   //lookat
               cam_lookup);   //lookup

    int i = int(gl_LocalInvocationID.x) / 2;
    int j = int(gl_LocalInvocationID.x) % 2;
    
    vec2 offset = vec2(1.0 * i / NS_X, 1.0 * j / NS_Y) - vec2(0.5);
    vec2 shootat = convertCoord(texCoord + offset, NX, NY);

    Ray ray = getRay(cam, shootat);
    vec3 texel = color(ray, reflect_type);
    // vec3 texel = fuck(ray, map_id);
    // texel = sqrt(texel);

    saveTex(texel, texCoord, map_id);
}