/**********************************************************************
Copyright ©2015 Advanced Micro Devices, Inc. All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

•   Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
•   Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or
other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
********************************************************************/

#define PI 3.14159265358979323846f

typedef struct _ray
{
    float4 o;
    float4 d;
} ray;

typedef struct _Intersection
{
    float4 uvwt;
    int shapeid;
    int primid;
} Intersection;

typedef struct _Camera
{
    // Camera coordinate frame
    float3 forward;
    float3 right;
    float3 up;
    float3 p;

    // Image plane width & height in scene units
    float2 dim;

    // Near and far Z
    float2 zcap;
    float  fovy;
    float  aspect;
} Camera;

typedef struct _Rng
{
    uint val;
} Rng;

typedef struct _Material
{
    float4 kd;
    float4 ks;
    float3 ke;
    int kdmapidx;
    int nmapidx;
    float metallic;
    float subsurface;
    float specular;
    float roughness;
    float speculartint;
    float anisotropic;
    float sheen;
    float sheentint;
    float clearcoat;
    float clearcoatgloss;
} Material;

enum TextureFormat
{
    UNKNOWN,
    RGBA8,
    RGBA16,
    RGBA32
};

// Texture description
typedef struct _Texture
{
    // Texture width, height and depth
    int w;
    int h;
    int d;
    // Offset in texture data array
    int dataoffset;
    int fmt;
    int extra;
} Texture;

// Shape description
typedef struct _Shape
{
    // Shape starting index in indices_ array
    int startidx;
    // Number of primitives in the shape
    int numprims;
    // Start vertex
    int startvtx;
    // Number of vertices
    int numvertices;
	// Linear motion vector
	float3 linearvelocity;
	// Angular velocity
	float4 angularvelocity;
	// Transform
	float4 m0;
	float4 m1;
	float4 m2;
	float4 m3;
} Shape;

// Emissive object
typedef struct _Emissive
{
    // Shape index
    int shapeidx;
    // Polygon index
    int primidx;
} Emissive;

// Unitility functions
#ifndef APPLE

float4 make_float4(float x, float y, float z, float w)
{
    float4 res;
    res.x = x;
    res.y = y;
    res.z = z;
    res.w = w;
    return res;
}

float3 make_float3(float x, float y, float z)
{
    float3 res;
    res.x = x;
    res.y = y;
    res.z = z;
    return res;
}

float2 make_float2(float x, float y)
{
    float2 res;
    res.x = x;
    res.y = y;
    return res;
}

int2 make_int2(int x, int y)
{
    int2 res;
    res.x = x;
    res.y = y;
    return res;
}

#endif

float3 transform_point(float3 p, float4 m0, float4 m1, float4 m2, float4 m3)
{
	float3 res;
	res.x = m0.s0 * p.x + m0.s1 * p.y + m0.s2 * p.z + m0.s3;
	res.y = m1.s0 * p.x + m1.s1 * p.y + m1.s2 * p.z + m1.s3;
	res.z = m2.s0 * p.x + m2.s1 * p.y + m2.s2 * p.z + m2.s3;
	return res;
}

float3 transform_vector(float3 p, float4 m0, float4 m1, float4 m2, float4 m3)
{
	float3 res;
	res.x = m0.s0 * p.x + m0.s1 * p.y + m0.s2 * p.z;
	res.y = m1.s0 * p.x + m1.s1 * p.y + m1.s2 * p.z;
	res.z = m2.s0 * p.x + m2.s1 * p.y + m2.s2 * p.z;
	return res;
}

float4 quaternion_mul(float4 q1, float4 q2)
{
	float4 res;
	res.x = q1.y*q2.z - q1.z*q2.y + q2.w*q1.x + q1.w*q2.x;
	res.y = q1.z*q2.x - q1.x*q2.z + q2.w*q1.y + q1.w*q2.y;
	res.z = q1.x*q2.y - q2.x*q1.y + q2.w*q1.z + q1.w*q2.z;
	res.w = q1.w*q2.w - q1.x*q2.x - q1.y*q2.y - q1.z*q2.z;
	return res;
}

float4 quaternion_conjugate(float4 q)
{
	return make_float4(-q.x, -q.y, -q.z, q.w);
}

float4 quaternion_inverse(float4 q)
{
	float sqnorm = q.x*q.x + q.y*q.y + q.z*q.z + q.w*q.w;

	if (sqnorm != 0.f)
	{
		return quaternion_conjugate(q) / sqnorm;
	}
	else
	{
		return make_float4(0.f, 0.f, 0.f, 1.f);
	}
}

float3 rotate_vector(float3 v, float4 q)
{
	float4 qinv = quaternion_inverse(q);
	float4 vv = make_float4(v.x, v.y, v.z, 0);
	return quaternion_mul(q, quaternion_mul(vv, qinv)).xyz;
}

uint WangHash(uint seed)
{
    seed = (seed ^ 61) ^ (seed >> 16);
    seed *= 9;
    seed = seed ^ (seed >> 4);
    seed *= 0x27d4eb2d;
    seed = seed ^ (seed >> 15);
    return seed;
}

uint RandUint(Rng* rng)
{
    rng->val = WangHash(1664525U * rng->val + 1013904223U);
    return rng->val;
}

float RandFloat(Rng* rng)
{
    return ((float)RandUint(rng)) / 0xffffffffU;
}

void InitRng(uint seed, Rng* rng)
{
    rng->val = WangHash(seed);
    for (int i=0;i<100;++i)
        RandFloat(rng);
}

float3 lerp(float3 a, float3 b, float w)
{
    return a + w*(b-a);
}

void CartesianToSpherical ( float3 cart, float* r, float* phi, float* theta )
{
    float temp = atan2(cart.x, cart.z);
    *r = sqrt(cart.x*cart.x + cart.y*cart.y + cart.z*cart.z);
    *phi = (float)((temp >= 0)?temp:(temp + 2*PI));
    *theta = acos(cart.y/ *r);
}

/// Uniformly sample a triangle
void SampleTriangleUniform(
    // RNG to use
    Rng* rng,
    // Triangle vertices
    float3 v0,
    float3 v1,
    float3 v2,
    // Triangle normals
    float3 n0,
    float3 n1,
    float3 n2,
    // Point and normal
    float3* p,
    float3* n,
    // PDF
    float* pdf
    )
{
    // Sample 2D value
    float a = RandFloat(rng);
    float b = RandFloat(rng);

    // Calculate point
    *p = (1.f - sqrt(a)) * v0.xyz + sqrt(a) * (1.f - b) * v1.xyz + sqrt(a) * b * v2.xyz;
    // Calculate normal
    *n = normalize((1.f - sqrt(a)) * n0.xyz + sqrt(a) * (1.f - b) * n1.xyz + sqrt(a) * b * n2.xyz);
    // PDF
    *pdf = 1.f / (length(cross((v2-v0), (v1-v0))) * 0.5f);
}

/// Get vector orthogonal to a given one
float3 GetOrthoVector(float3 n)
{
    float3 p;
    if (fabs(n.z) > 0.707106781186547524401f) {
        float k = sqrt(n.y*n.y + n.z*n.z);
        p.x = 0; p.y = -n.z/k; p.z = n.y/k;
    }
    else {
        float k = sqrt(n.x*n.x + n.y*n.y);
        p.x = -n.y/k; p.y = n.x/k; p.z = 0;
    }

    return p;
}

///< Power heuristic for multiple importance sampling
float PowerHeuristic(int nf, float fpdf, int ng, float gpdf)
{
    float f = nf * fpdf;
    float g = ng * gpdf;
    return (f*f) / (f*f + g*g);
}

///< Sample hemisphere with cos weight
float3 SampleHemisphere(
    // RNG to use
    Rng* rng,
    // Hemisphere normal
    float3 n,
    // Cos power
    float e)
{
    // Construct basis
    float3 u = GetOrthoVector(n);
    float3 v = cross(u, n);
    u = cross(n, v);

    // Calculate 2D sample
    float r1 = RandFloat(rng);
    float r2 = RandFloat(rng);

    // Transform to spherical coordinates
    float sinpsi = sin(2*PI*r1);
    float cospsi = cos(2*PI*r1);
    float costheta = pow(1.f - r2, 1.f/(e + 1.f));
    float sintheta = sqrt(1.f - costheta * costheta);

    // Return the result
    return normalize(u * sintheta * cospsi + v * sintheta * sinpsi + n * costheta);
}

/// Sample 2D texture described by texture in texturedata pool
float3 Sample2D(Texture const* texture, __global char const* texturedata, float2 uv)
{
    // Get width and height
    int width = texture->w;
    int height = texture->h;

    // Find the origin of the data in the pool
    __global char const* mydata = texturedata + texture->dataoffset;

    // Handle wrap
    uv -= floor(uv);

    // Reverse Y
    // it is needed as textures are loaded with Y axis going top to down
    // and our axis goes from down to top
    uv.y = 1.f - uv.y;

    // Figure out integer coordinates
    int x = floor(uv.x * width);
    int y = floor(uv.y * height);

    // Calculate samples for linear filtering
    int x1 = min(x + 1, width - 1);
    int y1 = min(y + 1, height - 1);

    // Calculate weights for linear filtering
    float wx = uv.x * width - floor(uv.x * width);
    float wy = uv.y * height - floor(uv.y * height);

    if (texture->fmt == RGBA32)
    {
        __global float3 const* mydataf = (__global float3 const*)mydata;

        // Get 4 values
        float3 valx = *(mydataf + width * y + x);
        float3 valx1 = *(mydataf + width * y + x1);
        float3 valy1 = *(mydataf + width * y1 + x);
        float3 valx1y1 = *(mydataf + width * y1 + x1);

        // Filter and return the result
        return lerp(lerp(valx, valx1, wx), lerp(valy1, valx1y1, wx), wy);
    }
    else if (texture->fmt == RGBA16)
    {
		__global half const* mydatah = (__global half const*)mydata;

        // Get 4 values
        float3 valx = vload_half4(width * y + x, mydatah).xyz;
		float3 valx1 = vload_half4(width * y + x1, mydatah).xyz;
		float3 valy1 = vload_half4(width * y1 + x, mydatah).xyz;
		float3 valx1y1 = vload_half4(width * y1 + x1, mydatah).xyz;

        // Filter and return the result
        return lerp(lerp(valx, valx1, wx), lerp(valy1, valx1y1, wx), wy);
    }
	else
	{
		__global uchar4 const* mydatac = (__global uchar4 const*)mydata;

        // Get 4 values
        uchar4 valx = *(mydatac + width * y + x);
        uchar4 valx1 = *(mydatac + width * y + x1);
        uchar4 valy1 = *(mydatac + width * y1 + x);
        uchar4 valx1y1 = *(mydatac + width * y1 + x1);

        float3 valxf = make_float3((float)valx.x / 255.f, (float)valx.y / 255.f, (float)valx.z / 255.f);
        float3 valx1f = make_float3((float)valx1.x / 255.f, (float)valx1.y / 255.f, (float)valx1.z / 255.f);
        float3 valy1f = make_float3((float)valy1.x / 255.f, (float)valy1.y / 255.f, (float)valy1.z / 255.f);
        float3 valx1y1f = make_float3((float)valx1y1.x / 255.f, (float)valx1y1.y / 255.f, (float)valx1y1.z / 255.f);

        // Filter and return the result
        return lerp(lerp(valxf, valx1f, wx), lerp(valy1f, valx1y1f, wx), wy);
	}
}

/// Sample 2D lat-long IBL map
float3 SampleEnvMap(Texture const* texture, __global char const* texturedata, float3 d)
{
    // Transform to spherical coords
    float r, phi, theta;
    CartesianToSpherical(d, &r, &phi, &theta);

    // Map to [0,1]x[0,1] range and reverse Y axis
    float2 uv;
    uv.x = phi / (2*PI);
    uv.y = 1.f - theta / PI;

    // Sample the texture
    return Sample2D(texture, texturedata, uv);
}

/// Get material data from hardcoded value or texture
float3 GetValue(
    // Value
    float3 v,
    // Texture index (can be -1 whihc means no texture)
    int texidx, 
    // Texture coordinate
    float2 uv,
    // Texture data
    __global Texture const* textures,
    __global char const* texturedata)
{
    // If texture present sample from texture
    if (texidx != -1)
    {
        // Get texture desc
        Texture texture = textures[texidx];
        // Sample diffuse texture
        return Sample2D(&texture, texturedata, uv);
    }

    // Return fixed color otherwise
    return v;
}


/// Construct tangent basis on the fly and apply normal map
float3 ApplyNormalMap(Texture* normalmap, __global char const* texturedata, float2 uv, float3 n, float3 v0, float3 v1, float3 v2, float2 uv0, float2 uv1, float2 uv2)
{
    /// From PBRT book
    float du1 = uv0.x - uv2.x;
    float du2 = uv1.x - uv2.x;
    float dv1 = uv0.y - uv2.y;
    float dv2 = uv1.y - uv2.y;

    float3 dp1 = v0 - v2;
    float3 dp2 = v1 - v2;

    float det = du1 * dv2 - dv1 * du2;


    float3 dpdu, dpdv;

    if (det != 0.f)
    {
        float invdet = 1.f / det;
        dpdu = normalize(( dv2 * dp1 - dv1 * dp2) * invdet);
        dpdv = normalize((-du2 * dp1 + du1 * dp2) * invdet);
    }
    else
    {
        dpdu = GetOrthoVector(n);
        dpdv = cross(n, dpdu);
    }

    float3 newnormal = cross(dpdv, dpdu);

    // Now n, dpdu, dpdv is orthonormal basis
    float3 mappednormal = 2.f * Sample2D(normalmap, texturedata, uv) - make_float3(1.f, 1.f, 1.f);

    // Return mapped version
    return normalize(mappednormal.z *  n + mappednormal.x * dpdu + mappednormal.y * dpdv);  
}

///< Shlick Fresnel term approx for Disney BRDF
float SchlickFresnel(const float u)
{
    const float m = clamp(1.f - u, 0.f, 1.f);
    const float m2 = m * m;
    return m2 * m2 * m;
}

float SchlickFresnelEta(float eta, float cosi)
{
    const float f = ((1.f - eta) / (1.f + eta)) * ((1.f - eta) / (1.f + eta));
    const float m = 1.f - fabs(cosi);
    const float m2 = m*m;
    return f + (1.f - f) * m2 * m2 * m;
}

float Gtr1(const float ndoth, const float a) 
{
    if (a >= 1.f)
        return 1.f / PI;

    const float a2 = a * a;
    const float t = 1.f + (a2 - 1.f) * ndoth * ndoth;

    return (a2 - 1.f) / (PI * log(a2) * t);
}

float Gtr2Aniso(
    const float ndoth, 
    const float hdotx, 
    const float hdoty,
    const float ax, 
    const float ay
    ) 
{
    return clamp(
        1.f / (PI * ax * ay * 
        ((hdotx / ax)*(hdotx / ax) + (hdoty / ay)*(hdoty / ay) + ndoth * ndoth) * 
        ((hdotx / ax)*(hdotx / ax) + (hdoty / ay)*(hdoty / ay) + ndoth * ndoth)),
        0.f, 10.f); // A trick to avoid fireflies
}

float SmithGggx(const float ndotv, const float alphag) 
{
    float a = alphag * alphag;
    float b = ndotv * ndotv;
    return 1.f / (ndotv + sqrt(a + b - a * b));
}

///< Evaluate actual layers
float3 DisneyEvaluateLayers(
    Material *material,
    const float3 n, 
    const float3 v, 
    const float3 l, 
    const float3 h,
    const float ds, 
    const float2 uv,
    __global Texture const* textures,
    __global char const* texturedata
    )
{
    // Read material parameters
    const float3 basecolor = GetValue(material->kd.xyz, material->kdmapidx, uv, textures, texturedata);
    const float metallic = material->metallic;
    const float subsurface = material->subsurface;
    const float specular = material->specular;
    const float roughness = material->roughness;
    const float speculartint = material->speculartint;
    const float sheen = material->sheen;
    const float sheentint = material->sheentint;
    const float clearcoat = material->clearcoat;
    const float clearcoatgloss = material->clearcoatgloss;

    const float ndotv = dot(n, v);
    const float ndotl = dot(n, l);

    const float ndoth = dot(n, h);
    const float ldoth = dot(l, h);

    const float3 cdlin = basecolor;
    const float cdlum = .3f * cdlin.x + .6f * cdlin.y + .1f * cdlin.z; // luminance approx.

    const float3 ctint = (cdlum > 0.f) ? (cdlin / cdlum) : make_float3(1.f,1.f,1.f); // normalize lum. to isolate hue+sat
    const float3 cspec0 = mix(specular * .08f * mix(1.f, ctint, speculartint), cdlin, metallic);
    const float3 csheen = mix(1.f, ctint, sheentint);

    // Diffuse fresnel - go from 1 at normal incidence to .5 at grazing
    // and mix in diffuse retro-reflection based on roughness
    const float fl = SchlickFresnel(ndotl);
    const float fv = SchlickFresnel(ndotv);
    const float fd90 = .5f + 2.f * ldoth * ldoth * roughness;
    const float fd = mix(1.f, fd90, fl) * mix(1.f, fd90, fv);

    // Based on Hanrahan-Krueger brdf approximation of isotropic bssrdf
    // 1.25 scale is used to (roughly) preserve albedo
    // fss90 used to flatten retroreflection based on roughness
    const float fss90 = ldoth * ldoth * roughness;
    const float fss = mix(1.f, fss90, fl) * mix(1.f, fss90, fv);
    const float ss = 1.25f * (fss * (1.f / (ndotl + ndotv) - .5f) + .5f);

    // Specular
    const float fh = SchlickFresnel(ndotl);
    const float3 fs = mix(cspec0, 1.f, fh);
    const float roughg = (roughness * .5f + .5f) * (roughness * .5f + .5f);
    const float gs = SmithGggx(ndotl, roughg) * SmithGggx(ndotv, roughg);

    // Sheen
    const float3 fsheen = fh * sheen * csheen;

    // Clearcoat (ior = 1.5 -> f0 = 0.04)
    const float dr = Gtr1(ndoth, mix(.1f, .001f, clearcoatgloss));
    const float fr = mix(.04f, 1.f, fh);
    const float gr = SmithGggx(ndotl, .25f) * SmithGggx(ndotv, .25f);

    float f = 1.f;
    const float et = material->kd.w;
    if (et != 0.f) 
    {
        const float3 wo = -v;
        const float cosi = dot(n, wo);
        f = 1.f - SchlickFresnelEta(et, cosi);
    }

    return f * (
        // Diffuse layer
        ((1.f / PI) * mix(fd, ss, subsurface) * cdlin + fsheen) * (1.f - metallic)
        +
        // 1st Specular lobe
        gs * fs * ds
        +
        // 2nd Specular lobe
        //
        clamp(clearcoat * gr * fr * dr, 0.f, .5f)
        );
}



///< Evaluate Disney principled BRDF
float3 DisneyEvaluate(Material *mat, 
    float3 n, 
    float3 wi,
    float3 wo, 
    float2 uv, 
    __global Texture const* textures, 
    __global char const* texturedata
) 
{
    const float3 v = wi;
    const float3 l = wo;

    if (dot(v, n) < 0.f)
        return 0.f;
    if (dot(l, n) <= 0.f)
        return 0.f;

    const float anisotropic = mat->anisotropic;
    const float roughness = mat->roughness;

    const float aspect = sqrt(1.f - anisotropic * .9f);
    const float ax = max(.001f, roughness * roughness / aspect);
    const float ay = max(.001f, roughness * roughness * aspect);

    float3 x, y;
    x = GetOrthoVector(n);
    y = normalize(cross(n, x));

    const float3 h = normalize(l + v);

    const float ndoth = dot(n, h);
    const float hdotx = dot(h, x);
    const float hdoty = dot(h, y);
    const float ds = Gtr2Aniso(ndoth, hdotx, hdoty, ax, ay);

    return DisneyEvaluateLayers(mat, n, v, l, h, ds, uv, textures, texturedata);
}

///< Sample Disney BRDF
float3 DisneySample(
    Rng* rng,
    Material *mat, 
    float3 n, 
    float3 wi,
    float2 uv, 
    __global Texture const* textures, 
    __global char const* texturedata,
    float3* wo,
    float* pdf
) 
{
    const float weight2 = mat->metallic;
    const float weight1 = 1.f - weight2;

    const float3 v = wi;
    if (dot(v, n) < 0.f)
        n = -n;

    const float anisotropic = mat->anisotropic;
    const float roughness = mat->roughness;

    const float aspect = sqrt(1.f - anisotropic * .9f);
    const float ax = max(.001f, roughness * roughness / aspect);
    const float ay = max(.001f, roughness * roughness * aspect);

    float3 x, y;
    x = GetOrthoVector(n);
    y = normalize(cross(n, x));

    const float u = RandFloat(rng);
    const bool samplediffuse = (u < weight1);

    float ds;
    float3 h, l;
    float diffusepdf, specularpdf;
    if (samplediffuse) 
    {
        // Sample the diffuse layer
        l = SampleHemisphere(rng, n, 0.f);

        h = normalize(l + v);

        const float ndoth = dot(n, h);
        const float hdotx = dot(h, x);
        const float hdoty = dot(h, y);
        ds = Gtr2Aniso(ndoth, hdotx, hdoty, ax, ay);

        const float ldoth = dot(l, h);
        diffusepdf = dot(l, n) / PI;
        specularpdf = (ds * ndoth) / (4.f * ldoth);
    }
    else {
        // Sample the specular layer
        const float u1 = RandFloat(rng);
        const float u2 = RandFloat(rng);

        h = normalize(sqrt(u2 / (1.f - u2)) * (ax * cos(2.f * PI * u1) + ay * sin(2.f * PI * u1)) + n);
        l = 2.f * dot(v, h) * h - v;

        const float ndoth = dot(n, h);
        const float hdotx = dot(h, x);
        const float hdoty = dot(h, y);
        ds = Gtr2Aniso(ndoth, hdotx, hdoty, ax, ay);

        const float ldoth = dot(l, h);
        if (ldoth < .01f) 
        {
            *pdf = 0.f;
            return make_float3(0.f, 0.f, 0.f);
        }

        diffusepdf = dot(l, n) / PI;
        specularpdf = (ds * ndoth) / (4.f * ldoth);

        if (specularpdf < 1.f) 
        {
            *pdf = 0.f;
            return make_float3(0.f, 0.f, 0.f);
        }
    }

    *wo = l;
    *pdf = weight1 * diffusepdf + weight2 * specularpdf;

    return DisneyEvaluateLayers(mat, n, v, l, h, ds, uv, textures, texturedata);
}

float DisneyGetPdf(
    Material *mat, 
    float3 n, 
    float3 wi, 
    float3 wo,
    float2 uv,
    __global Texture const* textures,
    __global char const* texturedata
    ) 
{
    const float3 v = wi;
    const float3 l = wo;

    if (dot(v, n) < 0.f)
        return 0.f;
    if (dot(l, n) <= 0.f)
        return 0.f;

    const float weight2 = mat->metallic;
    const float weight1 = 1.f - weight2;

    // Diffuse pdf
    const float diffusepdf = dot(l, n) / PI;

    // Specular pdf
    const float anisotropic = mat->anisotropic;
    const float roughness = mat->roughness;

    const float aspect = sqrt(1.f - anisotropic * .9f);
    const float ax = max(.001f, roughness * roughness / aspect);
    const float ay = max(.001f, roughness * roughness * aspect);

    float3 x, y;
    x = GetOrthoVector(n);
    y = normalize(cross(n, x));

    const float3 h = normalize(l + v);

    const float ndoth = dot(n, h);
    const float hdotx = dot(h, x);
    const float hdoty = dot(h, y);
    const float ldoth = dot(l, h);
    const float ds = Gtr2Aniso(ndoth, hdotx, hdoty, ax, ay);

    const float specularpdf = (ds * ndoth) / (4.f * ldoth);

    return weight1 * diffusepdf + weight2 * specularpdf;
}


///< Evaluate Disney principled BRDF
float3 LambertEvaluate(
    Material *mat,
    float3 n,
    float3 wi,
    float3 wo,
    float2 uv,
    __global Texture const* textures,
    __global char const* texturedata
    )
{
    const float3 kd = GetValue(mat->kd.xyz, mat->kdmapidx, uv, textures, texturedata);

    return kd * max(0.f, dot(wo, n)) / PI;
}

float LambertGetPdf(
    Material *mat,
    float3 n,
    float3 wi,
    float3 wo,
    float2 uv,
    __global Texture const* textures,
    __global char const* texturedata
    )
{
    return max(0.f, dot(n, wo)) / PI;
}

///< Sample Disney BRDF
float3 LambertSample(
    Rng* rng,
    Material *mat,
    float3 n,
    float3 wi,
    float2 uv,
    __global Texture const* textures,
    __global char const* texturedata,
    float3* wo,
    float* pdf
    )
{
    const float3 kd = GetValue(mat->kd.xyz, mat->kdmapidx, uv, textures, texturedata);

    *wo = SampleHemisphere(rng, n, 1.f);

    *pdf = max(0.f, dot(n, *wo)) / PI;

    return kd * (*pdf);
}

///< Evaluate material at shading point
float3 EvaluateBxdf(Material* mat, float3 wi, float3 wo, float3 n, float2 uv, __global Texture const* textures, __global char const* texturedata)
{
#ifdef USE_DISNEY_BRDF
    return DisneyEvaluate(mat, n, wi, wo, uv, textures, texturedata);
#else
    return LambertEvaluate(mat, n, wi, wo, uv, textures, texturedata);
#endif
}

float3 SampleBxdf(Rng* rng, Material* mat, float3 wi, float3 n, float2 uv, __global Texture const* textures, __global char const* texturedata, float3* wo, float* pdf)
{
#ifdef USE_DISNEY_BRDF
    return DisneySample(rng, mat, n, wi, uv, textures, texturedata, wo, pdf);
#else
    return LambertSample(rng, mat, n, wi, uv, textures, texturedata, wo, pdf);
#endif
}

float GetPdfBxfd(Material* mat, float3 wi, float3 wo, float3 n, float2 uv, __global Texture const* textures, __global char const* texturedata)
{
#ifdef USE_DISNEY_BRDF
    return DisneyGetPdf(mat, n, wi, wo, uv, textures, texturedata);
#else
    return LambertGetPdf(mat, n, wi, wo, uv, textures, texturedata);
#endif
}

///< Light sampling 
float3 EvaluateLight(float3 p, float3 n, float3 wo, int envmapidx, float envmapmul, __global Texture const* textures, __global char const* texturedata)
{
    // Sample envmap
    Texture envtexture = textures[envmapidx];
    return envmapmul * SampleEnvMap(&envtexture, texturedata, wo);
}

float3 SampleLight(Rng* rng, float3 p, float3 n, int envmapidx, float envmapmul, __global Texture const* textures, __global char const* texturedata, float3* wo, float* pdf)
{
    // Generate direction
    *wo = SampleHemisphere(rng, n, 1.f);

    // Envmap PDF 
    *pdf = dot(n, *wo) / PI;

    // Sample envmap
    Texture envtexture = textures[envmapidx];
    return envmapmul * SampleEnvMap(&envtexture, texturedata, *wo);
}

float GetPdfLight(float3 p, float3 n, float3 wo)
{
    return dot(n, wo) > 0 ? (dot(n, wo) / PI) : 0.f;
}


/// Ray generation kernel
/// Rays are generated from camera position to viewing plane
/// using random sample distribution within the pixel
///
__kernel void GenerateRays2D(
    // Camera buffer
    __global Camera const* camera,
    // Image resolution
    int imgwidth,
    int imgheight,
    // RNG seed value
    int randseed,
    // Output rays
    __global ray* rays)
{
    int2 global_id;
    global_id.x  = get_global_id(0);
    global_id.y  = get_global_id(1);

    // Handle only working subset
    if (global_id.x < imgwidth && global_id.y < imgheight)
    {
        __global ray* myray = rays + global_id.y * imgwidth + global_id.x;

        Rng rng;
        InitRng(randseed +  global_id.x * 157 - 33 * global_id.y, &rng);

        // Calculate [0..1] image plane sample
        float2 imgsample;
        imgsample.x = (float)global_id.x / imgwidth + RandFloat(&rng) / imgwidth;
        imgsample.y = (float)global_id.y / imgheight + RandFloat(&rng) / imgheight;

        // Transform into [-0.5,0.5]
        float2 hsample = imgsample - make_float2(0.5f, 0.5f);
        // Transform into [-dim/2, dim/2]
        float2 csample = hsample * camera->dim;

        // Direction to image plane
        myray->d.xyz = normalize(camera->zcap.x * camera->forward + csample.x * camera->right + csample.y * camera->up);
        // Origin == camera position
        myray->o.xyz = camera->p + camera->zcap.x * myray->d.xyz;
        // Zcap == (znear, zfar)
        myray->o.w = camera->zcap.y;
		// Generate random time from 0 to 1
        myray->d.w = RandFloat(&rng);
    }
}

///< Sample environemnt light and produce light samples and shadow rays into separate buffer
__kernel void SampleEnvironment(
    // Ray batch
    __global ray const* rays,
    // Intersection data
    __global Intersection const* isects,
    // Hit indices
    __global int const* hitindices,
    // Pixel indices
    __global int const* pixelindices,
    // Number of rays
    int numhits,
    // Vertices
    __global float3 const* vertices,
    // Normals
    __global float3 const* normals,
    // UVs
    __global float2 const* uvs,
    // Indices
    __global int const* indices,
    // Shapes
    __global Shape const* shapes,
    // Material IDs
    __global int const* materialids,
    // Materials
    __global Material const* materials,
    // Textures
    __global Texture const* textures,
    // Texture data
    __global char const* texturedata,
    // Environment texture index
    int envmapidx,
    // Envmap multiplier
    float envmapmul,
    // Emissives
    __global Emissive const* emissives,
    // Number of emissive objects
    int numemissives,
    // RNG seed
    int rngseed,
    // Shadow rays
    __global ray* shadowrays,
    // Light samples
    __global float3* lightsamples,
    // Number of shadow rays to spawn
    int numlightsamples,
    // Path throughput
    __global float3* throughput,
    // Indirect rays
    __global ray* indirectrays
    )
{
    int globalid = get_global_id(0);

    if (globalid < numhits)
    {
        // Fetch index
        int hitidx = hitindices[globalid];
        int pixelidx = pixelindices[globalid];

        // Fetch incoming ray
        float3 wi = -normalize(rays[hitidx].d.xyz);
		float time = rays[hitidx].d.w;

        // Determine shape and polygon
        int shapeid = isects[hitidx].shapeid - 1;
        int primid = isects[hitidx].primid;
        float2 uv = isects[hitidx].uvwt.xy;

        // Extract shape data
        Shape shape = shapes[shapeid];
		float3 linearvelocity = shape.linearvelocity;
		float4 angularvelocity = shape.angularvelocity;

        // Fetch indices starting from startidx and offset by primid
        int i0 = indices[shape.startidx + 3*primid];
        int i1 = indices[shape.startidx + 3*primid + 1];
        int i2 = indices[shape.startidx + 3*primid + 2];

        // Fetch normals
        float3 n0 = normals[shape.startvtx + i0];
        float3 n1 = normals[shape.startvtx + i1];
        float3 n2 = normals[shape.startvtx + i2];

        // Fetch positions
        float3 v0 = vertices[shape.startvtx + i0];
        float3 v1 = vertices[shape.startvtx + i1];
        float3 v2 = vertices[shape.startvtx + i2];

        // Fetch UVs
        float2 uv0 = uvs[shape.startvtx + i0];
        float2 uv1 = uvs[shape.startvtx + i1];
        float2 uv2 = uvs[shape.startvtx + i2];

        // Calculate barycentric position and normal
        float3 n = normalize((1.f - uv.x - uv.y) * n0 + uv.x * n1 + uv.y * n2);
        float3 v = (1.f - uv.x - uv.y) * v0 + uv.x * v1 + uv.y * v2;
        float2 t = (1.f - uv.x - uv.y) * uv0 + uv.x * uv1 + uv.y * uv2;

        // Get material at shading point
        int matidx = materialids[shape.startidx / 3 + primid];
        Material mat = materials[matidx];

        // Apply transform & linear motion blur
        v = rotate_vector(v, angularvelocity);
        v = transform_point(v, shape.m0, shape.m1, shape.m2, shape.m3);
        v += (linearvelocity * time);
        // MT^-1 should be used if scale is present
        n = rotate_vector(n, angularvelocity);
        n = transform_vector(n, shape.m0, shape.m1, shape.m2, shape.m3);

        // Apply normal map if needed
        if (mat.nmapidx != -1)
        {
            Texture normalmap = textures[mat.nmapidx];
            n = ApplyNormalMap(&normalmap, texturedata, t, n, v0, v1, v2, uv0, uv1, uv2);
        }

        if (dot(wi, n) < 0.f)
        {
            n = -n;
        }

        // Prepare RNG for light sampling
        Rng rng;
        InitRng(rngseed + globalid * 157, &rng);

        float lightpdf = 0.f;
        float bxdfpdf = 0.f;
        float3 wo;
        float3 le;

        // Generate shadow rays and light samples
        for (int i=0; i<numlightsamples/2; ++i)
        {
            // STEP1: Sample light and apply MIS

            // Generate direction
            le = SampleLight(&rng, v, n, envmapidx, envmapmul, textures, texturedata, &wo, &lightpdf);
            bxdfpdf = GetPdfBxfd(&mat, wi, wo, n, uv, textures, texturedata);
            float lightsampling_weight = PowerHeuristic(1, lightpdf, 1, bxdfpdf);

            float3 radiance = 0.f;

            if (lightpdf > 0.f)
            {
                // Generate light sample
                radiance = lightsampling_weight * le * EvaluateBxdf(&mat, wi, wo, n, t, textures, texturedata) * throughput[pixelidx] * max(0.f, dot(n, wo)) / lightpdf;
            }

            // Generate ray
            shadowrays[globalid * numlightsamples + 2*i].d.xyz = wo;
            shadowrays[globalid * numlightsamples + 2*i].o.xyz = v + 0.0001f * wo;
            shadowrays[globalid * numlightsamples + 2*i].o.w = 100000000.f;
            shadowrays[globalid * numlightsamples + 2*i].d.w = rays[hitidx].d.w;
            lightsamples[globalid * numlightsamples + 2 * i] = clamp(radiance, 0.f, 10.f);

            radiance = 0.f;
            {
                // Sample BxDF
                float3 bxdf = SampleBxdf(&rng, &mat, wi, n, t, textures, texturedata, &wo, &bxdfpdf);
                lightpdf = GetPdfLight(v, n, wo);
                float bxdfsampling_weight = PowerHeuristic(1, bxdfpdf, 1, lightpdf);
                le = EvaluateLight(v, n, wo, envmapidx, envmapmul, textures, texturedata); 

                if (bxdfpdf > 0.f)
                {
                    radiance = bxdfsampling_weight * le * bxdf * throughput[pixelidx] * max(0.f, dot(n, wo)) / bxdfpdf;
                }
            }

            // Generate ray
            shadowrays[globalid * numlightsamples + 2 * i + 1].d.xyz = wo;
            shadowrays[globalid * numlightsamples + 2 * i + 1].o.xyz = v + 0.0001f * wo;
            shadowrays[globalid * numlightsamples + 2 * i + 1].o.w = 100000000.f;
            shadowrays[globalid * numlightsamples + 2 * i + 1].d.w = rays[hitidx].d.w;
            lightsamples[globalid * numlightsamples + 2 * i + 1] = clamp(radiance, 0.f, 10.f);
        }

        // Generate indirect ray and update throughput
        float pdf;
        float3 bxdf = SampleBxdf(&rng, &mat, wi, n, t, textures, texturedata, &wo, &pdf);

        if (pdf > 0.001f && dot(n, wo) > 0.001f)
            throughput[pixelidx] *= (bxdf * max(0.f, dot(n, wo)) / pdf);
        else
            throughput[pixelidx] = 0.f;

        indirectrays[globalid].d.xyz = normalize(wo);
        indirectrays[globalid].o.xyz = v + 0.001f * normalize(wo);
        indirectrays[globalid].o.w = 100000000.f;
        indirectrays[globalid].d.w = rays[hitidx].d.w;
    }
}



///< Sample environemnt light and produce light samples and shadow rays into separate buffer
__kernel void SampleOcclusion(
    // Ray batch
    __global ray const* rays,
    // Intersection data
    __global Intersection const* isects,
    // Hit indices
    __global int const* hitindices,
    // Pixel indices
    __global int const* pixelindices,
    // Number of rays
    int numhits,
    // Vertices
    __global float3 const* vertices,
    // Normals
    __global float3 const* normals,
    // UVs
    __global float2 const* uvs,
    // Indices
    __global int const* indices,
    // Shapes
    __global Shape const* shapes,
    // Material IDs
    __global int const* materialids,
    // Materials
    __global Material const* materials,
    // Textures
    __global Texture const* textures,
    // Texture data
    __global char const* texturedata,
    // Environment texture index
    int envmapidx,
    // Envmap multiplier
    float envmapmul,
    // Emissives
    __global Emissive const* emissives,
    // Number of emissive objects
    int numemissives,
    // RNG seed
    int rngseed,
    // Shadow rays
    __global ray* shadowrays,
    // Light samples
    __global float3* lightsamples,
    // Number of shadow rays to spawn
    int numlightsamples,
    // Path throughput
    __global float3* throughput,
    // Indirect rays
    __global ray* indirectrays
    )
{
    int globalid = get_global_id(0);

    if (globalid < numhits)
    {
        // Fetch index
        int hitidx = hitindices[globalid];
        int pixelidx = pixelindices[globalid];

        // Fetch incoming ray
        float3 wi = -normalize(rays[hitidx].d.xyz);
        float time = rays[hitidx].d.w;

        // Determine shape and polygon
        int shapeid = isects[hitidx].shapeid - 1;
        int primid = isects[hitidx].primid;
        float2 uv = isects[hitidx].uvwt.xy;

        // Extract shape data
        Shape shape = shapes[shapeid];
        float3 linearvelocity = shape.linearvelocity;
        float4 angularvelocity = shape.angularvelocity;

        // Fetch indices starting from startidx and offset by primid
        int i0 = indices[shape.startidx + 3 * primid];
        int i1 = indices[shape.startidx + 3 * primid + 1];
        int i2 = indices[shape.startidx + 3 * primid + 2];

        // Fetch normals
        float3 n0 = normals[shape.startvtx + i0];
        float3 n1 = normals[shape.startvtx + i1];
        float3 n2 = normals[shape.startvtx + i2];

        // Fetch positions
        float3 v0 = vertices[shape.startvtx + i0];
        float3 v1 = vertices[shape.startvtx + i1];
        float3 v2 = vertices[shape.startvtx + i2];

        // Fetch UVs
        float2 uv0 = uvs[shape.startvtx + i0];
        float2 uv1 = uvs[shape.startvtx + i1];
        float2 uv2 = uvs[shape.startvtx + i2];

        // Calculate barycentric position and normal
        float3 n = normalize((1.f - uv.x - uv.y) * n0 + uv.x * n1 + uv.y * n2);
        float3 v = (1.f - uv.x - uv.y) * v0 + uv.x * v1 + uv.y * v2;
        float2 t = (1.f - uv.x - uv.y) * uv0 + uv.x * uv1 + uv.y * uv2;

        n = transform_vector(n, shape.m0, shape.m1, shape.m2, shape.m3);

        if (dot(wi, n) < 0.f)
        {
            n = -n;
        }

        // Prepare RNG for light sampling
        Rng rng;
        InitRng(rngseed + globalid * 157, &rng);

        float3 wo;
        float3 le;

        // Generate shadow rays and light samples
        for (int i = 0; i<numlightsamples / 2; ++i)
        {
            // Generate direction
            wo = SampleHemisphere(&rng, n, 1.f);

            // Generate ray
            shadowrays[globalid * numlightsamples + 2 * i].d.xyz = wo;
            shadowrays[globalid * numlightsamples + 2 * i].o.xyz = v + 0.001f * wo;
            shadowrays[globalid * numlightsamples + 2 * i].o.w = envmapmul;
            shadowrays[globalid * numlightsamples + 2 * i].d.w = rays[hitidx].d.w;

            // Generate ray
            shadowrays[globalid * numlightsamples + 2 * i + 1].d.xyz = wo;
            shadowrays[globalid * numlightsamples + 2 * i + 1].o.xyz = v + 0.001f * wo;
            shadowrays[globalid * numlightsamples + 2 * i + 1].o.w = envmapmul;
            shadowrays[globalid * numlightsamples + 2 * i + 1].d.w = rays[hitidx].d.w;
        }
    }
}

///< Illuminate missing rays
__kernel void ShadeMiss(
    // Ray batch
    __global ray const* rays,
    // Intersection data
    __global Intersection const* isects,
    // Pixel indices
    __global int const* pixelindices,
    // Number of rays
    int numrays,
    // Textures
    __global Texture const* textures,
    // Texture data
    __global char const* texturedata,
    // Environment texture index
    int envmapidx,
    // Througput
    __global float3 const* througput,
    // Output values
    __global float4* output
    )
{
    int globalid = get_global_id(0);

    if (globalid < numrays)
    {
        int pixelidx = pixelindices[globalid];

        // In case of a miss
        if (isects[globalid].shapeid < 0)
        {
            // Sample IBL for that ray
            Texture texture = textures[envmapidx];
            // Multiply by througput
            output[pixelidx].xyz += SampleEnvMap(&texture, texturedata, rays[globalid].d.xyz) * througput[pixelidx];
        }

        output[pixelidx].w += 1.f;
    }
}

///< Handle light samples and visibility info and add contribution to final buffer
__kernel void GatherLightSamples(
    // Pixel indices
    __global int const* pixelindices,
    // Number of rays
    int numrays,
    // Number of light samples per ray
    int numlightsamples,
    // Shadow rays hits
    __global int const* shadowhits,
    // Light samples
    __global float3 const* lightsamples,
    // Througput
    __global float3 const* throughput,
    // Radiance sample buffer
    __global float4* output
    )
{
    int globalid = get_global_id(0);

    if (globalid < numrays)
    {
        // Prepare accumulator variable
        float3 radiance = make_float3(0.f, 0.f, 0.f);

        // Get pixel id for this sample set
        int pixelidx = pixelindices[globalid];

        // Start collecting samples
        for (int i=0; i<numlightsamples; ++i)
        {
            // If shadow ray didn't hit anything and reached skydome
            if (shadowhits[globalid * numlightsamples + i] == -1)
            {
                // Add its contribution to radiance accumulator
                radiance += lightsamples[globalid * numlightsamples + i];
            }
        }

        // Divide by number of light samples (samples already have built-in throughput
        output[pixelidx].xyz += (radiance / numlightsamples);
    }
}


///< Handle light samples and visibility info and add contribution to final buffer
__kernel void GatherOcclusion(
    // Pixel indices
    __global int const* pixelindices,
    // Number of rays
    int numrays,
    // Number of light samples per ray
    int numlightsamples,
    // Shadow rays hits
    __global int const* shadowhits,
    // Light samples
    __global float3 const* lightsamples,
    // Througput
    __global float3 const* throughput,
    // Radiance sample buffer
    __global float4* output
    )
{
    int globalid = get_global_id(0);

    if (globalid < numrays)
    {
        // Prepare accumulator variable
        float occlusion = 0.f;

        // Get pixel id for this sample set
        int pixelidx = pixelindices[globalid];

        // Start collecting samples
        for (int i = 0; i<numlightsamples; ++i)
        {
            // If shadow ray didn't hit anything and reached skydome
            if (shadowhits[globalid * numlightsamples + i] == 1)
            {
                occlusion += 1.f;
            }
        }

        float visibility = 1.f - occlusion / numlightsamples;

        // Divide by number of light samples (samples already have built-in throughput
        output[pixelidx].xyz += make_float3(visibility, visibility, visibility);
        output[pixelidx].w += 1.f;
    }
}


///< Restore pixel indices after compaction
__kernel void RestorePixelIndices(
    // Compacted indices
    __global int const* compacted_indices,
    // Number of compacted indices
    int numitems,
    // Previous pixel indices
    __global int const* previndices,
    // New pixel indices
    __global int* newindices
    )
{
    int globalid = get_global_id(0);

    // Handle only working subset
    if (globalid < numitems)
    {
        newindices[globalid] = previndices[compacted_indices[globalid]];
    }
}

///< Restore pixel indices after compaction
__kernel void ConvertToPredicate(
    // Intersections
    __global Intersection const* isects,
    // Number of compacted indices
    int numitems,
    // Previous pixel indices
    __global int* predicate
    )
{
    int globalid = get_global_id(0);

    // Handle only working subset
    if (globalid < numitems)
    {
        predicate[globalid] = isects[globalid].shapeid >= 0 ? 1 : 0;
    }
}


