#include "/lib/block_id.glsl"


float jitteredHash(inout float seed) {
    float jitterSize = 1.0/8.0;//4 stratum
    float uni = hash1(seed);
    float jitter = hash1(seed)*jitterSize - jitterSize*0.5;
    return mod(uni + jitter, 1.0);
}

const float piOver4 = 0.785398163397448;
const float piOver2 = 1.570796326794;

vec2 concentricSampleDisk(vec2 random) {
    random = 2.0 * random - 1.0;

    if (all(equal(random, vec2(0)))) return vec2(0);

    float theta, r;
    if (abs(random.x) > abs(random.y)) {
        r = random.x;
        theta = piOver4 * (random.y / random.x);
    } else {
        r = random.y;
        theta = piOver2 - piOver4 * (random.x / random.y);
    }
    return r * vec2(cos(theta), sin(theta));
}
vec3 cosWeightedRandomHemisphereDirection(vec3 n, vec2 r) {
    vec2 d = concentricSampleDisk(r);
    float z = sqrt(max(0.0, 1.0 - d.x * d.x - d.y * d.y));

    vec3  uu = normalize(cross(n, vec3(0.0, 1.0, 1.0)));
    vec3  vv = cross(uu, n);
    vec3  rr = vec3(d.x*uu + d.y*vv + z*n);
    return normalize(rr);
}

float GGX_D(float dotNH, float alpha2) {
    float den = (alpha2 - 1.0) * dotNH * dotNH + 1.0;
    return alpha2 / (PI * den * den);
}

//BRDF
vec3 ggx_light(vec3 n, vec3 v, vec3 l, vec2 specularity, vec3 albedo) {
    float alpha2 = specularity.x * specularity.x;

    float dotNL = clamp(dot(n, l), 0., 1.);
    float dotNV = clamp(dot(n, v), 0., 1.);

    vec3 h = normalize(v + l);
    float dotNH = clamp(dot(n, h), 0., 1.);
    float dotLH = clamp(dot(l, h), 0., 1.);

    // GGX microfacet distribution function
    float D = GGX_D(dotNH, alpha2);

    // Fresnel with Schlick approximation
    vec3 F0 = specularity.y*albedo;

    vec3 F = F0 + (1.0 - F0) * pow(1. - dotLH, 5.0);

    // Smith joint masking-shadowing function
    float k = (specularity.x + 1)*(specularity.x + 1)/8.0;
    float G = 1.0 / ((dotNL * (1.0 - k) + k) * (dotNV * (1.0 - k) + k));

    //pdf = max(D * dotNH /(4.0*dotLH), 0.00001);
    return max(D * F * G, 0.0001);
}

vec3 ggx(vec3 n, vec3 v, vec3 l, vec2 specularity, vec3 albedo, out float fresnel) {
    float alpha2 = specularity.x * specularity.x;

    float dotNL = clamp(dot(n, l), 0., 1.);
    float dotNV = clamp(dot(n, v), 0., 1.);

    vec3 h = normalize(v + l);
    float dotNH = clamp(dot(n, h), 0., 1.);
    float dotLH = clamp(dot(l, h), 0., 1.);

    // GGX microfacet distribution function
    float D = GGX_D(dotNH, alpha2);

    // Fresnel with Schlick approximation
    vec3 F0 = specularity.y*albedo;

    vec3 F = F0 + (1.0 - F0) * pow(1. - dotLH, 5.0);
    fresnel = F.r;

    // Smith joint masking-shadowing function
    float k = .5 * alpha2;
    float G = 1.0 / ((dotNL * (1.0 - k) + k) * (dotNV * (1.0 - k) + k));

    //pdf = max(D * dotNH /(4.0*dotLH), 0.00001);
    return max(D * F * G, 0.0001);
}

//Sample ggx direction & PDF
vec3 ggx_sample(vec3 n, vec3 v, float roughness, out float pdf, vec2 rand) {

    float alpha2 = roughness * roughness;

    float epsilon = clamp(clamp(rand.x, 0.0, 0.5), 0.001, 1.);
    float cosTheta2 = (1. - epsilon) / (epsilon * (alpha2 - 1.) + 1.);
    float cosTheta = sqrt(cosTheta2);
    float sinTheta = sqrt(1. - cosTheta2);

    float phi = 2. * PI * rand.y;

    // Spherical to cartesian
    vec3 t = normalize(cross(n.yzx, n));
    vec3 b = cross(n, t);

    vec3 microNormal = (t * cos(phi) + b * sin(phi)) * sinTheta + n * cosTheta;

    vec3 l = reflect(-v, microNormal);

    // Sample weight
    float den = (alpha2 - 1.) * cosTheta2 + 1.;
    float D = alpha2 / (PI * den * den);
    float p = D * cosTheta / (4. * dot(microNormal, v));
    pdf = (.5 / PI) / (p + 1e-6);
    pdf = 1.0/max(pdf, 0.00001);

//    if (dot(l, n) < 0.)
//        weight = 0.000000;

    return l;
}

vec3 getGGXDominantDirection(vec3 normal, vec3 view, float roughness) {
    float f = (1.0 - roughness) * (sqrt(1.0 - roughness) + roughness);
    vec3 R = reflect(-view, normal);
    return normalize(mix(normal, R, f));
}

// vec3 ggx_vndf_sampler(vec3 normal, vec3 V, float roughness, vec2 U) {
//     //Tail clamp
//     U.x = U.x < 0.2 ? U.x : 0.2;

//     // Rotate to tangent space
//     V = Rotate(V, normal, vec3(0, 1, 0));

//     // stretch view
//     vec3 v = normalize(vec3(V.x * roughness, V.y, V.z * roughness));

//     // orthonomal basis
//     vec3 t1 = (v.y < 0.999f) ? normalize(cross(v, vec3(0, 1, 0))) : vec3(1, 0, 0);
//     vec3 t2 = cross(t1, v);

//     // sample point with polar coordinates (r, phi)
//     float a = 1.0f / (1.0f + v.y);
//     float r = sqrt(U.x);
//     float phi = (U.y < a) ? (U.y / a) * PI 
//                          : PI + (U.y - a) / (1.0f - a) * PI;
//     float p1 = r * cos(phi);
//     float p2 = r * sin(phi) * ((U.y < a) ? 1.0f : v.y);

//     // compute normal
//     vec3 n = p1 * t1 + p2 * t2 + sqrt(max(0.0f, 1.0f - p1 * p1 - p2 * p2)) * v;

//     // unstretch
//     n = normalize(vec3(roughness * n.x, max(0.0f, n.y), roughness * n.z));

//     // Reflect
//     vec3 reflected = reflect(-V, n);



vec3 parseSpecular(vec3 data) {
    float roughness = clamp(1 - data.x, 0.089, 1.0);
        roughness = roughness * roughness;
    float f0 = clamp(data.y, 0.02, 0.9);

    float emission = data.b;
    // return vec3(1.0, 1.0, 0.0);
    return vec3(roughness, f0, emission);
}

// float antiAliasSpecular(in vec3 normal, in float roughness2) {
//     vec3 dndu = dFdx( normal ), dndv = dFdy( normal );
// 	float variance = 3004.0 * ( dot( dndu, dndu ) + dot( dndv, dndv ) );
// 	float kernelRoughness2 = min(variance, 1.0 );
// 	return sqrt(clamp( roughness2 + kernelRoughness2 , 0.0, 1.0));
// }

float schlick(vec3 v, vec3 n, vec3 l, float F0) {
    vec3 h = normalize(v + l);
    float dotLH = clamp(dot(l, h), 0.00, 1.);

    return clamp01(F0 + (1.0 - F0) * pow(1. - dotLH, 5.0));
}

float schlick(vec3 incident, vec3 normal, float F0) {
    float cosTheta = clamp01(dot(-incident, normal));

    return clamp(F0 + (1.0 - F0) * pow(1.0 - cosTheta, 5.0), 0.0, 1.0);
}

vec3 schlick3(vec3 incident, vec3 normal, vec3 F0) {
    float cosTheta = clamp01(dot(-incident, normal));

    return clamp01(F0 + (1.0 - F0) * pow(1.0 - cosTheta, 5.0));
}

float schlickInverse(float f0, float VoH) {
    return 1.0 - clamp01(f0 + (1.0 - f0) * pow(1.0 - VoH, 5.0));
}

float diffuseHammon(vec3 normal, vec3 viewDir, vec3 lightDir, float roughness) {
    float nDotL     = max0(dot(normal, lightDir));

    if (nDotL <= 0.0) return (0.0);

    float nDotV     = max0(dot(normal, viewDir));

    float lDotV     = max0(dot(lightDir, viewDir));

    vec3 halfWay    = normalize(viewDir + lightDir);
    float nDotH     = max0(dot(normal, halfWay));

    float facing    = lDotV * 0.5 + 0.5;

    float singleRough = facing * (0.9 - 0.4 * facing) * ((0.5 + nDotH) * rcp(max(nDotH, 0.02)));
    float singleSmooth = 1.05 * schlickInverse(0.0, nDotL) * schlickInverse(0.0, max0(nDotV));

    float single    = clamp01(mix(singleSmooth, singleRough, roughness) * rcp(PI));
    float multi     = 0.1159 * roughness;

    return clamp01((multi + single) * nDotL);
}

bool isMetal(float f0) {
   return (f0 > 220.0/255.0);
}

vec3 computeNEEWeight(vec3 normal, vec3 incident, vec3 sunDir, vec3 albedo, vec2 specularity, bool isSpecular) {
    vec3 BRDF;
    float isMetal = float(isMetal(specularity.y));
    float cosTheta = clamp01(dot(normal, sunDir));
    if(!isSpecular) {
        BRDF = albedo * diffuseHammon(normal, -incident, sunDir, specularity.r);
        // BRDF = albedo * rcp(PI) * cosTheta * 1.5;
    } else {
        BRDF = mix(vec3(1), albedo, isMetal) * ggx_light(normal, -incident, sunDir, specularity, albedo) * cosTheta;
    }
    
    return BRDF;
}

vec3 computeWeight(vec3 normal, vec3 incident, vec3 reflectedDir, vec3 albedo, vec2 specularity, bool isSpecular, bool doAlbedoTint, out float fresnel) {
    float cosTheta = clamp01(dot(normal, reflectedDir));
    vec3 BRDF;
    bool isMetal = isMetal(specularity.y);
    if(!isSpecular) {
        // BRDF = mix(vec3(1), albedo, float(doAlbedoTint)) * diffuseHammon(normal, -incident, reflectedDir, specularity.r);
        BRDF = mix(vec3(1), albedo, float(doAlbedoTint)) * rcp(PI) * cosTheta;
    } else {
        BRDF = mix(vec3(1), albedo, float(isMetal && doAlbedoTint)) * ggx(normal, -incident, reflectedDir, specularity, albedo, fresnel) * cosTheta;
    }
    return BRDF;
}

// vec3 reflectRay(in vec3 normal, in vec3 incident, in vec3 albedo, in vec2 specularity, in bool isSpecular, bool doAlbedoTint, out vec3 weight, inout float seed) {
    
//     vec2 r = hash2(seed);
//     // r.x = jitteredHash(seed);
//     // r.y = jitteredHash(seed);

//     vec3 reflectedDir;
//     float pdf;
//     if(!isSpecular) {
//         reflectedDir = cosWeightedRandomHemisphereDirection(normal, r);
//         float cosTheta = clamp01(dot(normal,  reflectedDir));
//         pdf = max(cosTheta*rcp(PI), 0.00001); 
//     } else {
//         reflectedDir = ggx_sample(normal, -incident, specularity.r, pdf, r);
//         pdf = max(pdf, 0.00001); 
//     }

//     weight = computeWeight(normal, incident, reflectedDir, albedo, specularity, isSpecular, doAlbedoTint) * rcp(pdf);

//     return reflectedDir;
// }

vec3 reflectDiffuseRay(in vec3 normal, in vec3 incident, in vec3 albedo, in vec2 specularity, bool doAlbedoTint, out vec3 weight, in vec2 r) {
    

    vec3 reflectedDir;
    float pdf;
    reflectedDir = cosWeightedRandomHemisphereDirection(normal, r);
    float cosTheta = clamp01(dot(normal, reflectedDir));
    pdf = max(cosTheta*rcp(PI), 0.00001); 
    float f;
    weight = computeWeight(normal, incident, reflectedDir, albedo, specularity, false, doAlbedoTint, f) * rcp(pdf);

    return reflectedDir;
}

vec3 reflectDiffuseRay(in vec3 normal, in vec3 incident, in vec3 albedo, in vec2 specularity, bool doAlbedoTint, out vec3 weight, inout float seed) {
    
    vec2 r = hash2(seed);

    return reflectDiffuseRay(normal, incident, albedo, specularity, doAlbedoTint, weight, r);
}

vec3 reflectSpecRay(in vec3 normal, in vec3 incident, in vec3 albedo, in vec2 specularity, bool doAlbedoTint, out vec4 weight_invpdf, vec2 r, out float fresnel) {
    // r.x = mix(0.0, r.x, 0.5);
    vec3 reflectedDir;
    float pdf;
    reflectedDir = ggx_sample(normal, -incident, specularity.r, pdf, r);
    pdf = max(pdf, 0.00001); 

    weight_invpdf = vec4(computeWeight(normal, incident, reflectedDir, albedo, specularity, true, doAlbedoTint, fresnel), rcp(pdf));

    return reflectedDir;
}

vec3 DFGPolynomial( vec3 F0, float roughness, float ndotv ) {
    float gloss = 1.0 - roughness;
    float x = gloss;
    float y = ndotv;
 
    float b1 = -0.1688;
    float b2 = 1.895;
    float b3 = 0.9903;
    float b4 = -4.853;
    float b5 = 8.404;
    float b6 = -5.069;
    float bias = clamp01( min( b1 * x + b2 * x * x, b3 + b4 * y + b5 * y * y + b6 * y * y * y ) );
 
    float d0 = 0.6045;
    float d1 = 1.699;
    float d2 = -0.5228;
    float d3 = -3.603;
    float d4 = 1.404;
    float d5 = 0.1939;
    float d6 = 2.661;
    float delta = clamp01( d0 + d1 * x + d2 * y + d3 * x * x + d4 * x * y + d5 * y * y + d6 * x * x * x );
    float scale = delta - bias;
 
    bias *= clamp01( 14.0 * dot(F0, vec3(0.5)) );
    return F0 * scale + bias;
}