// Copyright Epic Games, Inc. All Rights Reserved.

/*=============================================================================================
PathTracingSolidGlass.ush: Microfacet Refractive BSDF for solid glass
===============================================================================================*/

#pragma once

#include "PathTracingMaterialCommon.ush"
#include "PathTracingFresnel.ush"
#include "PathTracingGlossy.ush"
#include "PathTracingEnergyConservation.ush"

struct FRoughGlassData
{
	float3x3 Basis;
	float2 Alpha;
	float3 V;

	// glass lobe data
	float Eta;
	float F0;
	float EGlass;
};

FRoughGlassData PrepareRoughGlassData(FPathTracingPayload Payload, float3 V_World)
{
	FRoughGlassData Data = (FRoughGlassData)0;

	Data.Basis = GetGGXBasis(Payload.Roughness, Payload.Anisotropy, Payload.WorldNormal, Payload.WorldTangent, Data.Alpha);
	Data.V = mul(Data.Basis, V_World);
	const float NoV = saturate(Data.V.z);


	// NOTE: IsFrontFace() determines polygon orientation, because the normal is always flipped towards in the incoming ray
	// TODO: Maintain a refraction stack on the path tracing payload
	Data.Eta = Payload.IsFrontFace() ? Payload.Ior : rcp(Payload.Ior);
	Data.F0 = F0RGBToF0(Payload.SpecularColor);

	// correct for energy loss by scaling the whole BSDF
	Data.EGlass = GGXEnergyLookup(Payload.Roughness, NoV, Data.Eta);

	return Data;
}


FMaterialSample RoughGlass_SampleMaterial(
	float3 V_World,
	FPathTracingPayload Payload,
	float3 RandSample)
{
	const FRoughGlassData Data = PrepareRoughGlassData(Payload, V_World);
	const float3 V = Data.V;
	const float NoV = saturate(V.z);
	if (NoV == 0)
	{
		// invalid grazing angle
		return NullMaterialSample();
	}
	const float3 H = ImportanceSampleVisibleGGX(RandSample.xy, Data.Alpha, V, false).xyz;
	// Glass lobe (reflection and refraction)
	float3 L = 0;
	float F = 0;

	const bool bRefract = SampleRefraction(-V, H, Data.Eta, Data.F0, RandSample.z, L, F);
	// transform to world space
	const float3 L_World = normalize(mul(L, Data.Basis));

	const float2 GGXResult = bRefract ? GGXEvalRefraction(L, V, H, Data.Alpha, Data.Eta) : GGXEvalReflection(L, V, H, Data.Alpha, false);

	return CreateMaterialSample(L_World, Payload.BSDFOpacity * GGXResult.x / Data.EGlass, F * GGXResult.y, bRefract ? -1.0 : 1.0, Payload.Roughness, bRefract ? PATHTRACER_SCATTER_REFRACT : PATHTRACER_SCATTER_SPECULAR);
}

FMaterialEval RoughGlass_EvalMaterial(
	float3 V_World,
	float3 L_World,
	FPathTracingPayload Payload,
	float2 DiffuseSpecularScale
)
{
	const FRoughGlassData Data = PrepareRoughGlassData(Payload, V_World);
	const float3 V = Data.V;
	const float NoV = saturate(V.z);
	if (NoV == 0)
	{
		// invalid grazing angle
		return NullMaterialEval();
	}


	// move vectors into right shading frame
	const float3 L = mul(Data.Basis, L_World);

	FMaterialEval Result = NullMaterialEval();

	if (L.z > 0)
	{
		// reflection side
		const float3 H = normalize(V + L);

		const float NoL = saturate(L.z);
		const float VoH = saturate(dot(V, H));
		const float NoH = saturate(H.z);

		// Glass lobe
		const float2 GGXResult = GGXEvalReflection(L, V, H, Data.Alpha, false);
		const float Fg = FresnelReflectance(VoH, Data.Eta, Data.F0);
		const float GlassWeight = GGXResult.x * Payload.BSDFOpacity / Data.EGlass;
		const float GlassPdf = GGXResult.y * Fg;
		Result.AddLobeWithMIS(GlassWeight * DiffuseSpecularScale.y, GlassPdf, 1.0);
	}
	else if (L.z < 0)
	{
		// refraction side
		const float NoL = saturate(-L.z);
		float3 Ht = -(Data.Eta * L + V);
		Ht = normalize((Data.Eta < 1.0f) ? -Ht : Ht);

		const float VoH = dot(V, Ht);
		const float Fg = 1.0f - FresnelReflectance(VoH, Data.Eta, Data.F0);
		if (Fg > 0)
		{
			const float2 GGXResult = GGXEvalRefraction(L, V, Ht, Data.Alpha, Data.Eta);

			const float GlassWeight = GGXResult.x * Payload.BSDFOpacity / Data.EGlass;
			const float GlassPdf = GGXResult.y * Fg;

			Result.AddLobeWithMIS(GlassWeight * DiffuseSpecularScale.y, GlassPdf, 1.0);
		}
	}
	return Result;
}