vec3 tone(vec3 color, float gamma) //Reinhard based tone mapping, from: https://www.shadertoy.com/view/XdSSRw
{
	float white = 2.;
	float luma = dot(color, vec3(0.2126, 0.7152, 0.0722));
	float toneMappedLuma = luma * (1. + luma / (white*white)) / (1. + luma);
	color *= toneMappedLuma / luma;
	color = pow(color, vec3(1. / gamma));
	return color;
}

float rand(float co) { return fract(sin(co*(91.3458)) * 47453.5453); } //https://www.shadertoy.com/view/Xt23Ry
float rand(vec2 co){ return fract(sin(dot(co.xy ,vec2(12.9898,78.233))) * 43758.5453); }
//float rand(vec3 co){ return rand(co.xy+rand(co.z)); }

//Copyright (C) 2020 Nicolas Ortiz
// License Creative Commons Attribution-NonCommercial-ShareAlike 3.0 Unported License.

#define PI 3.1415926538f

//Primitives
bool trace_sphere (vec3 ro, vec3 rd, vec3 o, float r, float tmin, float tmax,
                  out vec3 p, out vec3 n, out vec3 t, out vec2 uv, out float dist)
{
	vec3 oc = ro - o;
    float a = dot (rd, rd);
    float b = dot (oc, rd);
    float c = dot (oc, oc) - r * r;
    float t0 = b * b - a * c;
    dist = tmax;

    if (t0 > 0.0)
    {
        float t1 = (-b - sqrt (t0)) / a;

        if (t1 < tmax && t1 > tmin)
        {
            dist = t1;
            p = ro + rd * dist;
            n = (p - o) * (1.0f / r);
            t = cross(vec3 (0, 1, 0), n);

            uv.x = (1.f + atan (n.z, n.x) / PI) * 0.5f;
            uv.y = acos (n.y) / PI;

            return true;
        }

        t1 = (-b + sqrt (t0)) / a;

        if (t1 < tmax && t1 > tmin)
        {
            dist = t1;
            p = ro + rd * dist;
            n = (p - o) * (1.0f / r);
            t = cross(vec3 (0, 1, 0), n);

            uv.x = (1.f + atan (n.z, n.x) /PI) * 0.5f;
            uv.y = acos (n.y) / PI;

            return true;
        }
    }

    return false;
}
//BSDFs/PDFs
vec3 sample_sphere (vec2 screen_uv) //Uniform sphere sample
{
    float cosPhi = 2.0 * rand (screen_uv*iTime) - 1.0;
    float sinPhi = sqrt (1.0 - cosPhi * cosPhi);
    float theta = 2.0 * PI * rand (rand (screen_uv*screen_uv * iTime));

    return vec3 (sinPhi * sin (theta),
        cosPhi,
        sinPhi * cos (theta));
}

vec3 sample_metal(vec3 n, vec3 rd, float fuzzy, vec2 screen_uv) //Mix (interpolate) reflection with diffuse by fuzz param basically
{
    vec3 reflected = normalize(rd - n * dot(n, rd) * 2.0f);
    return normalize(reflected + sample_sphere(screen_uv) * fuzzy);
}

vec3 sample_cosine_weighted(vec3 n, vec2 screen_uv) //cosine weighted hemisphere sampling
{
    float phi = 2.f * PI * rand (screen_uv);
    float r2 = rand(rand (screen_uv*iTime));
    float r2s = sqrt (r2);

    vec3 w = normalize (n);
    vec3 u = normalize(cross((abs (w.x) > .1 ? vec3 (0, 1, 0) : vec3 (1, 0, 0)), w));
    vec3 v = cross (w, u);

    return normalize(u * cos (phi) * r2s + v * sin (phi) * r2s + w * sqrt (1.f - r2));
}

vec3 sample_phong_metal(vec3 n, vec3 rd, float e, vec2 screen_uv) //phong for metals
{
    float phi = 2.f * PI * rand(screen_uv);
    float r2 = rand(rand(screen_uv));

    float cos_theta = pow(1.f - r2, 1.f / (e + 1.f));
    float sin_theta = sqrt(1.f - cos_theta * cos_theta);

    vec3 w = normalize((rd - n * dot(n, rd) * 2.0f));
    vec3 u = normalize(cross((abs(w.x) > 1.f ? vec3(0,1,0) : vec3(1,0,0)),w));
    vec3 v = cross(w, u);

    return normalize(u * cos(phi) * sin_theta + v * sin(phi) * sin_theta + w * cos_theta);
}
float F_Schlick (float cosine, float ref_idx) //schlick fresnel factor
{
    float r0 = (1.0f - ref_idx) / (1.0f + ref_idx);
    r0 = r0 * r0;
    return r0 + (1.0f - r0) * pow (1.0f - cosine, 5.0f);
}

bool do_refract(vec3 v, vec3 n, float ni_over_nt, out vec3 refr)
{
    vec3 uv = normalize(v);
    float dt = dot(uv, n);
    float t = 1.0f - ni_over_nt * ni_over_nt * (1.0f - dt * dt);

    if(t > 0.0f)
    {
        refr = (uv - n * dt) * ni_over_nt - n * sqrt(t);
     	return true;
    }

    return false;
}
vec3 sample_dielectric(vec3 n, vec3 rd, float ior, vec2 screen_uv)
{
    float idotn = dot (rd, n);

    vec3 outward_normal;
    float ni_over_nt;
    float cosine;
    if (idotn > 0.0f) //from inside to outside or the other way around (dot(raydir, normal))
    {
        outward_normal = -n;
        ni_over_nt = ior;
        cosine = idotn / length(rd);
        cosine = sqrt (1.0f - ior * ior * (1.0f - cosine * cosine));
    }
    else
    {
        outward_normal = n;
        ni_over_nt = 1.0f / ior;
        cosine = -idotn / length(rd);
    }

    vec3 refracted;
    float p;
    if (do_refract (rd, outward_normal, ni_over_nt, refracted)) //compute refr. dir and will it refract or reflect
        p = F_Schlick (cosine, ior); //probability of reflection is fresnel-schlick
    else
        p = 1.0f; //reflect 100%

    if (rand (screen_uv * iTime) < p)
        return normalize(normalize((rd - n * dot (n, rd) * 2.0))); //reflection
    else
        return normalize(refracted); //refraction
}


//Main
bool trace_scene(inout vec3 ro, inout vec3 rd, out vec3 d, out vec3 e, vec2 screen_uv)
{
    //n: normal, p: point, t: tangent, d: diffuse, e: emission, ro: ray origin, rd: ray direction, uv: tex coords
    vec3 n, p, t;
    vec2 uv;
    float tmax = 1000.f;
    float dist;
    //Here in the event of a ray-object collision we set diffuse and emission based on
    //e objects properties and set the direction of the ray based on the objects bsdf.
    //As well as tmax to the objects distance (doesnt make a difference yet since we lack depth sorting)

    bool o1 = trace_sphere(ro, rd, vec3(sin(iTime * 0.8f) - 0.8f,0.5f + sin(iTime * 2.f) * 0.5f,-10.f+sin(iTime * 1.5f)), 1.f, 0.001f, tmax, p, n, t, uv, dist);
    if(o1)
    {
        d = vec3(0.9,0.9,0.9);
        e = vec3(0);

        ro = p;
        rd = sample_dielectric(n, rd, 1.52f, screen_uv);//or 1.02 for less of an effect
        tmax = dist;
        return true;
    }

    bool o2 = trace_sphere(ro, rd, vec3(1.5f,0,-12.f + (sin((0.65f * PI) + iTime * 1.2f) * 1.5f)), 1.f, 0.001f, tmax, p, n, t, uv, dist);
    if(o2)
    {
        d = vec3(.7,.7,0);
        e = vec3(0);

        ro = p;
        rd = sample_metal(n, rd, 0.7f, screen_uv);
        tmax = dist;
        return true;
    }


    bool o3 = trace_sphere(ro, rd, vec3(-1.5f + sin(iTime*3.0f)*0.6f,1.f + sin(iTime*3.f)*0.6f,-12), 1.f, 0.001f, tmax, p, n, t, uv, dist);
    if(o3)
    {
        d = vec3(0,0,0);
        e = vec3(1,1,0.4f) * 30.f;

        ro = p;
        rd = vec3(0);
        tmax = dist;
        return true;
    }

    //disabled to lower compilation time
    /*bool o4 = trace_sphere(ro, rd, vec3(-3.f + sin(iTime * 0.5f) * 0.5f,0,-12.f + sin(iTime * 0.5f) * .3f), 1.f, 0.001f, tmax, p, n, t, uv, dist);
    if(o4)
    {
        d = texture(iChannel1, uv).rgb;
        e = vec3(0);

        ro = p;
        rd = sample_metal(n, rd, 0.0f, screen_uv);
        tmax = dist;
        return true;
    }*/

    /*bool o5 = trace_sphere(ro, rd, vec3(3.5f,0,-12.f), 1.f, 0.001f, tmax, p, n, t, uv, dist);
    if(o5)
    {
        d = texture(iChannel1, uv).rgb;
        e = vec3(0);

        ro = p;
        rd = sample_cosine_weighted(n, screen_uv);
        tmax = dist;
        return true;
    }*/

    bool o6 = trace_sphere(ro, rd, vec3(-1.5f,-1001,-12), 1000.f, 0.001f, tmax, p, n, t, uv, dist);
    if(o6)
    {
        d = vec3(.5,.2,.2); //diffuse color
        e = vec3(0); //emission

        ro = p; //ray origin out
        rd = sample_cosine_weighted(n, uv); //ray dir out
        tmax = dist; //for depth testing
        return true;
    }


    return false;
}

vec3 radiance(in vec3 ro, in vec3 rd, vec2 uv)
{
    vec3 att = vec3(1);
    vec3 col;

    for(int i = 0; i < 15; i++) //15 max bounces
    {
        vec3 d, e;

        if(!trace_scene(ro, rd, d, e, uv))
        {
            vec4 hdri = texture(iChannel0, rd);
            col += att * hdri.rgb * hdri.a;
            break;
        }

        col += att * e;    //Emission
        att *= d;         //Diffuse color

        if(i > 7) //Russian roulette sampling
        {
            float p = max(att.x, max(att.y, att.z));
            if(rand(uv) > p)
                break;

            att /= p;
        }
    }

    return col;
}

void mainImage( out vec4 fragColor, in vec2 fragCoord )
{
    int spp = 100; //Camera params
    float aperture = 0.3;//clamp(sin(iTime * 0.7f) + 0.2f, 0.f, 1.f);
    float vfov = 30.f;
    vec3 pos = vec3(-0.7f,0,0);
    float fl = 10.0f;

    vec3 right = vec3(1,0,0); //Camera direction (use for rotation)
    vec3 up = vec3(0,1,0);
    vec3 fwd = vec3(0,0,1);

    float aspect = iResolution.x/iResolution.y; //Perspective calculations (frustum)
    float hh = tan((vfov * (PI / 180.0f)) / 2.0f);
    float hw = aspect * hh;
    vec3 ll = pos - right * hw * fl - up * hh * fl - fwd * fl;
    vec3 h = right * fl * 2.0f * hw;
    vec3 v = up * fl * 2.0f * hh;

    vec3 color = vec3(0,0,0);
    for(int i = 0; i < spp; i++) //Supersampling
    {
        vec2 uv_o = vec2(rand(fragCoord.xy * float(i)), rand(rand(fragCoord.xy)* float(i))); //Random offset
        vec2 uv = (fragCoord + uv_o)/iResolution.xy; //Normalized screen coordinates with offset

        float r = sqrt(rand(iTime * uv));  //Disc sampling (DoF)
        float theta = rand(rand((iTime * uv))) * 2.0f * PI;
        vec3 ds = vec3(cos(theta), sin(theta), 0) * (aperture/2.0f);
        vec3 o = right * ds.x + up * ds.y; //DoF offset

        vec3 ro = pos + o; //ray origin
        vec3 rd = ll + h * uv.x + v * uv.y - pos - o; //ray dir

        color += clamp(tone(radiance(ro, rd, uv), 1.0f), 0.f, 1.f);
    }

    fragColor = vec4(color / float(spp), 1); // Final color is average of samples tonemapped
}
