#version 330
/*
  Copyright(C) 2018-2025, tetraface Inc. All rights reserved.
  
  This is a modified implementation of glTF-Sample-Viewer.
  The original code is here:
    https://github.com/KhronosGroup/glTF-Sample-Viewer
*/
/*
                                 Apache License
                           Version 2.0, January 2004
                        http://www.apache.org/licenses/
*/

uniform mat4 WorldView;
uniform mat4 WorldViewProj;
uniform mat4 ModelMatrix;
int Type;
uniform vec4 BaseColor;
//uniform float Metallic;
uniform vec4 Emissive;
uniform float Roughness;
uniform float Ior;
uniform vec3 Scattering;
uniform float Scanisotropy;
uniform float TRDepth;
uniform vec3 AttenuationColor;
uniform int HasBumpMap;
uniform vec3 NormalMapFlip;
uniform int HasRoughnessMap;
uniform int HasEmissiveMap;
uniform int LightNum;
uniform vec4 LightPos[4];
uniform vec3 LightCol[4];
uniform vec3 LightSpc[4];
uniform vec4 BackFaceColor;
uniform vec3 GlobalAmbient;
uniform int AlphaBlendMode;
uniform int HasTextureMap;
uniform float EnvMapFactor;
uniform float EnvMapRotate;
uniform int EnvSpecularMapMipLevel;
uniform int TransmissionMapMipLevel;
uniform vec3 ViewDir;
uniform vec3 CameraPos;
uniform int ClipPlaneNum;
uniform vec4 ClipPlane[4];

uniform sampler2D ColorMapSampler;
uniform sampler2D BumpMapSampler;
uniform sampler2D RoughnessMapSampler;
uniform sampler2D EmissiveMapSampler;
uniform sampler2D EnvDiffuseMapSampler;
uniform sampler2D EnvSpecularMapSampler;
uniform sampler2D BrdfLUTMapSampler;
uniform sampler2D TransmissionBufferSampler;

in vec3 VSWorldPos;
in vec4 VSColor;
in vec3 VSNormal;
in vec2 VSTexCoord;
#if USE_TANGENT
in vec3 VSTangent;
in vec3 VSBinormal;
#endif
#if MULTIUV >= 1
in vec2 VSTexCoord1;
#endif
#if MULTIUV >= 2
in vec2 VSTexCoord2;
#endif
#if MULTIUV >= 3
in vec2 VSTexCoord3;
#endif

out vec4 OutColor;

const float PI = 3.141592653589793;

#define float2 vec2
#define float3 vec3
#define float4 vec4
float lerp(float a, float b, float t) { return mix(a,b,t); }
vec3 lerp(vec3 a, vec3 b, float t) { return mix(a,b,t); }
vec3 mul(vec3 v, mat3 m) { return m * v; }
float saturate(float v) { return clamp(v, 0, 1); }
vec2 saturate(vec2 v) { return clamp(v, 0, 1); }
vec3 saturate(vec3 v) { return clamp(v, 0, 1); }
vec4 saturate(vec4 v) { return clamp(v, 0, 1); }
float atan2(in float y, in float x) { return x == 0.0 ? sign(y)*PI/2 : atan(y, x); }

vec2 GetTexCoord(int channel)
{
#if MULTIUV >= 1
	if((channel & 0x30) == 0x10) return VSTexCoord1;
#endif
#if MULTIUV >= 2
	if((channel & 0x30) == 0x20) return VSTexCoord2;
#endif
#if MULTIUV >= 3
	if((channel & 0x30) == 0x30) return VSTexCoord3;
#endif
	return VSTexCoord;
}

const float c_MinRoughness = 0.04;
const float Metallic = 0;

struct BRDFInput
{
	float metallic;
	float perceptualRoughness;
	float alphaRoughness;
	float3 diffuseColor;
	float3 f0;
};

float max3(float3 v) { return max(max(v.x, v.y), v.z); }
float3 to3(float v) { return float3(v,v,v); }
float clampedDot(float3 x, float3 y) { return saturate(dot(x, y)); }

vec2 dirToUV(vec3 dir)
{
	return vec2(
		atan2(dir.z, dir.x) / (PI*2),
		acos(dir.y) / PI);
}

vec2 envmapUV(vec3 dir)
{
	return dirToUV(dir) + vec2(EnvMapRotate, 0.0);
}

float LINEARtoSRGB(float c)
{
    return (c < 0.0031308) ? c * 12.92 : 1.055 * pow(abs(c), 0.41666) - 0.055;
}

vec3 LINEARtoSRGB(vec3 srgbIn)
{
	return vec3(
		LINEARtoSRGB(srgbIn.x),
		LINEARtoSRGB(srgbIn.y),
		LINEARtoSRGB(srgbIn.z));
}

float SRGBtoLINEAR(float c)
{
    return (c < 0.04045) ? c*0.0773993808 : pow(c*0.9478672986+0.0521327014, 2.4);
}

float4 SRGBtoLINEAR(float4 srgbIn)
{
	float3 linOut = float3(
		SRGBtoLINEAR(srgbIn.x),
		SRGBtoLINEAR(srgbIn.y),
		SRGBtoLINEAR(srgbIn.z));
	return float4(linOut,srgbIn.w);
}

float4 SRGBtoLINEAR(float4 srgbIn, bool is_srgb)
{
	if(is_srgb){
		return float4(SRGBtoLINEAR(srgbIn).xyz,srgbIn.w);
	}else{
		return srgbIn;
	}
}

float3 F_Schlick(float3 f0, float3 f90, float VdotH)
{
	return f0 + (f90 - f0) * pow(clamp(1.0 - VdotH, 0.0, 1.0), 5.0);
}

// Smith Joint GGX
// Note: Vis = G / (4 * NdotL * NdotV)
// see Eric Heitz. 2014. Understanding the Masking-Shadowing Function in Microfacet-Based BRDFs. Journal of Computer Graphics Techniques, 3
// see Real-Time Rendering. Page 331 to 336.
// see https://google.github.io/filament/Filament.md.html#materialsystem/specularbrdf/geometricshadowing(specularg)
float V_GGX(float NdotL, float NdotV, float alphaRoughness)
{
	float alphaRoughnessSq = alphaRoughness * alphaRoughness;

	float GGXV = NdotL * sqrt(NdotV * NdotV * (1.0 - alphaRoughnessSq) + alphaRoughnessSq);
	float GGXL = NdotV * sqrt(NdotL * NdotL * (1.0 - alphaRoughnessSq) + alphaRoughnessSq);

	float GGX = GGXV + GGXL;
	if (GGX > 0.0)
		return 0.5 / GGX;
	else
		return 0.0;
}

// The following equation(s) model the distribution of microfacet normals across the area being drawn (aka D())
// Implementation from "Average Irregularity Representation of a Roughened Surface for Ray Reflection" by T. S. Trowbridge, and K. P. Reitz
// Follows the distribution function recommended in the SIGGRAPH 2013 course notes from EPIC Games [1], Equation 3.
float D_GGX(float NdotH, float alphaRoughness)
{
	float alphaRoughnessSq = alphaRoughness * alphaRoughness;
	float f = (NdotH * NdotH) * (alphaRoughnessSq - 1.0) + 1.0;
	return alphaRoughnessSq / (PI * f * f);
}

//  https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#acknowledgments AppendixB
float3 BRDF_specularGGX(float3 f0, float3 f90, float alphaRoughness, float specularWeight, float VdotH, float NdotL, float NdotV, float NdotH)
{
	float3 F = F_Schlick(f0, f90, VdotH);
	float Vis = V_GGX(NdotL, NdotV, alphaRoughness);
	float D = D_GGX(NdotH, alphaRoughness);

	return specularWeight * F * Vis * D;
}

//https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#acknowledgments AppendixB
float3 BRDF_lambertian(float3 f0, float3 f90, float3 diffuseColor, float specularWeight, float VdotH)
{
	// see https://seblagarde.wordpress.com/2012/01/08/pi-or-not-to-pi-in-game-lighting-equation/
	return (1.0 - specularWeight * F_Schlick(f0, f90, VdotH)) * (diffuseColor / PI);
}

float3 getIBLRadianceGGX(float3 N, float3 V, float roughness, float3 F0, float specularWeight)
{
    float NdotV = clampedDot(N, V);
    float lod = roughness * float(EnvSpecularMapMipLevel-1);
    float3 reflection = normalize(reflect(-V, N));

    float2 brdfSamplePoint = saturate(float2(NdotV, 1.0 - roughness));
    float2 f_ab = texture(BrdfLUTMapSampler, brdfSamplePoint).xy;
    float3 specularLight = textureLod(EnvSpecularMapSampler, envmapUV(reflection), lod).xyz;

    // see https://bruop.github.io/ibl/#single_scattering_results at Single Scattering Results
    // Roughness dependent fresnel, from Fdez-Aguera
    float3 Fr = max(to3(1.0 - roughness), F0) - F0;
    float3 k_S = F0 + Fr * pow(1.0 - NdotV, 5.0);
    float3 FssEss = k_S * f_ab.x + f_ab.y;

    return specularWeight * specularLight * FssEss;
}

// specularWeight is introduced with KHR_materials_specular
float3 getIBLRadianceLambertian(float3 n, float3 v, float roughness, float3 diffuseColor, float3 F0, float specularWeight)
{
	float NdotV = clampedDot(n, v);
	float2 brdfSamplePoint = clamp(float2(NdotV, 1.0 - roughness), 0, 1);
	float2 f_ab = texture(BrdfLUTMapSampler, brdfSamplePoint).xy;

	float3 irradiance = texture(EnvDiffuseMapSampler, envmapUV(n)).xyz;

	// see https://bruop.github.io/ibl/#single_scattering_results at Single Scattering Results
	// Roughness dependent fresnel, from Fdez-Aguera

	float3 Fr = max(to3(1.0 - roughness), F0) - F0;
	float3 k_S = F0 + Fr * pow(1.0 - NdotV, 5.0);
	float3 FssEss = specularWeight * k_S * f_ab.x + f_ab.y; // <--- GGX / specular light contribution (scale it down if the specularWeight is low)

	// Multiple scattering, from Fdez-Aguera
	float Ems = (1.0 - (f_ab.x + f_ab.y));
	float3 F_avg = specularWeight * (F0 + (1.0 - F0) / 21.0);
	float3 FmsEms = Ems * FssEss * F_avg / (1.0 - F_avg * Ems);
	float3 k_D = diffuseColor * (1.0 - FssEss + FmsEms); // we use +FmsEms as indicated by the formula in the blog post (might be a typo in the implementation)

	return (FmsEms + k_D) * irradiance;
}

// Compute attenuated light as it travels through a volume.
float3 applyVolumeAttenuation(float3 radiance, float transmissionDistance, float3 attenuationColor, float attenuationDistance)
{
	if (attenuationDistance == 0.0)
	{
		// Attenuation distance is +inf (which we indicate by zero), i.e. the transmitted color is not attenuated at all.
		return radiance;
	}
	else
	{
		// Compute light attenuation using Beer's law.
		float3 attenuationCoefficient = -log(attenuationColor) / attenuationDistance;
		float3 transmittance = exp(-attenuationCoefficient * transmissionDistance); // Beer's law
		return transmittance * radiance;
	}
}

float3 getVolumeTransmissionRay(float3 n, float3 v, float thickness, float ior, mat4 modelMatrix)
{
	// Direction of refracted light.
	float3 refractionVector = refract(-v, normalize(n), 1.0 / ior);

	// Compute rotation-independant scaling of the model matrix.
	float3 modelScale;
	modelScale.x = length(float3(modelMatrix[0].xyz));
	modelScale.y = length(float3(modelMatrix[1].xyz));
	modelScale.z = length(float3(modelMatrix[2].xyz));

	// The thickness is specified in local space.
	return normalize(refractionVector) * thickness * modelScale;
}

float applyIorToRoughness(float roughness, float ior)
{
	// Scale roughness with IOR so that an IOR of 1.0 results in no microfacet refraction and
	// an IOR of 1.5 results in the default amount of microfacet refraction.
	return roughness * clamp(ior * 2.0 - 2.0, 0.0, 1.0);
}

float3 getTransmissionSample(float2 fragCoord, float roughness, float ior)
{
	float framebufferLod = float(TransmissionMapMipLevel-1) * applyIorToRoughness(roughness, ior);
	float4 transmittedLight = textureLod(TransmissionBufferSampler, float2(fragCoord.x,1.0-fragCoord.y), framebufferLod);
	return SRGBtoLINEAR(transmittedLight).xyz;
}

float3 getIBLVolumeRefraction(float3 n, float3 v, float perceptualRoughness, float3 baseColor, float3 f0, float3 f90,
	float3 position, mat4 modelMatrix, float ior, float thickness, float3 attenuationColor, float attenuationDistance)
{
	float3 transmissionRay = getVolumeTransmissionRay(n, v, thickness, ior, modelMatrix);
	float3 refractedRayExit = position + transmissionRay;

	// Project refracted vector on the framebuffer, while mapping to normalized device coordinates.
	float4 ndcPos = WorldViewProj * vec4(refractedRayExit, 1.0);
	float2 refractionCoords = ndcPos.xy / ndcPos.w;
	refractionCoords.y = -refractionCoords.y;
	refractionCoords += 1.0;
	refractionCoords /= 2.0;

	// Sample framebuffer to get pixel the refracted ray hits.
	float3 transmittedLight = getTransmissionSample(refractionCoords, perceptualRoughness, ior);

	float3 attenuatedColor = applyVolumeAttenuation(transmittedLight, length(transmissionRay), attenuationColor, attenuationDistance);

	// Sample GGX LUT to get the specular component.
	float NdotV = clampedDot(n, v);
	float2 brdfSamplePoint = saturate(float2(NdotV, 1.0 - perceptualRoughness));
	float2 brdf = texture(BrdfLUTMapSampler, brdfSamplePoint).xy;
	float3 specularColor = f0 * brdf.x + f90 * brdf.y;

	return (1.0 - specularColor) * attenuatedColor * baseColor;
}

float3 getPunctualRadianceTransmission(float3 normal, float3 view, float3 pointToLight, float alphaRoughness,
    float3 f0, float3 f90, float3 baseColor, float ior)
{
	float transmissionRougness = applyIorToRoughness(alphaRoughness, ior);

	float3 n = normalize(normal);           // Outward direction of surface point
	float3 v = normalize(view);             // Direction from surface point to view
	float3 l = normalize(pointToLight);
	float3 l_mirror = normalize(l + 2.0*n*dot(-l, n));     // Mirror light reflection vector on surface
	float3 h = normalize(l_mirror + v);            // Halfway vector between transmission light vector and v

	float D = D_GGX(clamp(dot(n, h), 0.0, 1.0), transmissionRougness);
	float3 F = F_Schlick(f0, f90, clamp(dot(v, h), 0.0, 1.0));
	float Vis = V_GGX(clamp(dot(n, l_mirror), 0.0, 1.0), clamp(dot(n, v), 0.0, 1.0), transmissionRougness);

	// Transmission BTDF
	return (1.0 - F) * baseColor * D * Vis;
}

BRDFInput BRDFSetup(float3 baseColor)
{
	BRDFInput brdfInput;
    // Metallic and Roughness material properties are packed together
    // In glTF, these factors can be specified by fixed scalar values
    // or from a metallic-roughness map
	brdfInput.metallic = Metallic;
    brdfInput.perceptualRoughness = Roughness;
	if(HasRoughnessMap != 0){
		// Roughness is stored in the 'g' channel, metallic is stored in the 'b' channel.
		// This layout intentionally reserves the 'r' channel for (optional) occlusion map data
		float4 mrSample = texture(RoughnessMapSampler, GetTexCoord(HasRoughnessMap));
		brdfInput.perceptualRoughness *= mrSample.g;
		//brdfInput.metallic *= mrSample.b;
	}
    brdfInput.perceptualRoughness = clamp(brdfInput.perceptualRoughness, 0.0, 1.0);
    brdfInput.metallic = clamp(brdfInput.metallic, 0.0, 1.0);
    // Roughness is authored as perceptual roughness; as is convention,
    // convert to material roughness by squaring the perceptual roughness [2].
    brdfInput.alphaRoughness = clamp(brdfInput.perceptualRoughness * brdfInput.perceptualRoughness, c_MinRoughness * c_MinRoughness, 1.0);

	brdfInput.diffuseColor = baseColor * (1.0 - brdfInput.metallic);
	brdfInput.f0 = to3(pow((Ior - 1.0) /  (Ior + 1.0), 2.0));

	// Achromatic f0 based on IOR.
	brdfInput.f0 = lerp(brdfInput.f0, baseColor, brdfInput.metallic);

	return brdfInput;
}

void main(void)
{
	for(int i=0; i<ClipPlaneNum; i++){
		if(dot(ClipPlane[i].xyz, VSWorldPos) + ClipPlane[i].w < 0)
			discard;
	}
	vec3 V = normalize(CameraPos - VSWorldPos);
	vec4 base_color = SRGBtoLINEAR(texture(ColorMapSampler, GetTexCoord(HasTextureMap)), (HasTextureMap&3)==1) * VSColor;
	if(base_color.w < 0.01)
		discard;

	vec3 N;
#if USE_TANGENT
	if(HasBumpMap != 0){
		vec3 bump_col = texture(BumpMapSampler, GetTexCoord(HasBumpMap)).xyz * 2 - 1;
		//bump_col *= float3(NormalScale, NormalScale, 1.0);
		bump_col *= NormalMapFlip;
		mat3 mtx = mat3(VSTangent, VSBinormal, VSNormal);
		N = normalize(mul(bump_col, mtx));
	}else
#endif
	{
		N = normalize(VSNormal);
	}
	if(gl_FrontFacing){
		base_color.xyz *= BackFaceColor.xyz;
		N = -N;
	}
	BRDFInput brdfInput = BRDFSetup(base_color.xyz);

	float ior = Ior;
	float3 transmission = float3(0,0,0);
	float transmissionFactor = (AlphaBlendMode != 0) ? 1.0 : 0.0;
	float thickness = 0;
	float3 attenuationColor = float3(0,0,0);
	float attenuationDistance = 0;
	if(Type == 1){
		attenuationColor = AttenuationColor;
		attenuationDistance = TRDepth;
	}
	float3 f90 = float3(1,1,1);
	float specularWeight = 1;

	float3 b_diffuse = float3(0,0,0);
	float3 b_specular = float3(0,0,0);
	for(int i=0; i<4 && i<LightNum; i++){
		float3 L;
		if(LightPos[i].w == 0){
			L = normalize(LightPos[i].xyz - VSWorldPos.xyz); // point light
		}else{
			L = LightPos[i].xyz; // directional light
		}
		float3 litcol = LightCol[i];
		float3 diffuse = float3(0,0,0);
		float3 specular = float3(0,0,0);

		float3 h = normalize(L + V);          // Direction of the vector between l and v, called halfway vector
		float NdotL = clampedDot(N, L);
		float NdotV = clampedDot(N, V);
		float NdotH = clampedDot(N, h);
		float VdotH = clampedDot(V, h);
		if (NdotL > 0.0 || NdotV > 0.0)
		{
			diffuse = litcol * NdotL * BRDF_lambertian(brdfInput.f0, f90, brdfInput.diffuseColor, specularWeight, VdotH);
			specular = litcol * NdotL * BRDF_specularGGX(brdfInput.f0, f90, brdfInput.alphaRoughness, specularWeight, VdotH, NdotL, NdotV, NdotH);
		}

		// If the light ray travels through the geometry, use the point it exits the geometry again.
		// That will change the angle to the light source, if the material refracts the light ray.
		float3 transmissionRay = getVolumeTransmissionRay(N, V, thickness, ior, ModelMatrix);
		float3 tL = normalize(L - transmissionRay);
		float3 transmittedLight = litcol * getPunctualRadianceTransmission(N, V, tL, brdfInput.alphaRoughness, brdfInput.f0, f90, brdfInput.diffuseColor, ior);
		if(Type == 1)
			transmittedLight = applyVolumeAttenuation(transmittedLight, length(transmissionRay), attenuationColor, attenuationDistance);
		transmission += transmissionFactor * transmittedLight;

		b_diffuse += diffuse;
		b_specular += specular;
	}
	b_diffuse *= PI;
	b_specular *= PI;

#if ENVMAP
	b_specular += getIBLRadianceGGX(N, V, brdfInput.perceptualRoughness, brdfInput.f0, specularWeight) * EnvMapFactor;
	b_diffuse += getIBLRadianceLambertian(N, V, brdfInput.perceptualRoughness, brdfInput.diffuseColor, brdfInput.f0, specularWeight) * EnvMapFactor;
#else // No envmap
	b_diffuse += base_color.xyz * GlobalAmbient.xyz;
#endif

	transmission += transmissionFactor * getIBLVolumeRefraction(
		N, V,
		brdfInput.perceptualRoughness,
		brdfInput.diffuseColor, brdfInput.f0, f90,
		VSWorldPos, ModelMatrix,
		ior, thickness, attenuationColor, attenuationDistance);
	b_diffuse = lerp(b_diffuse, transmission, transmissionFactor);

	vec3 color = b_diffuse + b_specular;

	if(HasEmissiveMap != 0){
		vec4 emissive = SRGBtoLINEAR(texture(EmissiveMapSampler, GetTexCoord(HasEmissiveMap)), (HasEmissiveMap&3)==1);
		color += emissive.xyz * Emissive.xyz;
	}else{
		color += Emissive.xyz;
	}
	vec4 col;
	col.xyz = LINEARtoSRGB(saturate(color));
	col.w = saturate(base_color.w);
	OutColor = col;
}

