#include "Common/Color.hlsli"
#include "Common/FrameBuffer.hlsli"
#include "Common/SharedData.hlsli"

#if !defined(DYNAMIC_CUBEMAPS) && defined(IBL)
#	undef IBL
#endif

struct VS_INPUT
{
	float4 Position: POSITION0;
#if !defined(ENVCUBE)
	float4 Normal: NORMAL0;
#endif
	float4 TexCoord0: TEXCOORD0;
#if defined(ENVCUBE)
	float4
#else
	int4
#endif
		TexCoord1: TEXCOORD1;
};

struct VS_OUTPUT
{
	float4 Position: SV_POSITION0;
	float4 Color: COLOR0;
	float2 TexCoord0: TEXCOORD0;
#if defined(ENVCUBE)
	float4 PrecipitationOcclusionTexCoord: TEXCOORD1;
#endif
#if defined(ENVCUBE) && defined(RAIN) && defined(EFFECTS11)
	float2 RaindropData: TEXCOORD2;
#endif
};

#ifdef VSHADER
cbuffer PerTechnique : register(b0)
{
	float2 ScaleAdjust : packoffset(c0);
};

cbuffer PerGeometry : register(b2)
{
	row_major float4x4 WorldViewProj;  // 0
	row_major float4x4 WorldView;      // 4
#	if defined(ENVCUBE)
	row_major float4x4 PrecipitationOcclusionWorldViewProj;  // 8, 16
#	endif
	float4 fVars0;        // 8, 16 ENVCUBE 12, 20
	float4 fVars1;        // 9, 17 ENVCUBE 13, 21
	float4 fVars2;        // 10, 18 ENVCUBE 14, 22
	float4 fVars3;        // 11, 19 ENVCUBE 15, 23
	float4 fVars4;        // 12, 20 ENVCUBE 16, 24
	float4 Color1;        // 13, 21 ENVCUBE 17, 25
	float4 Color2;        // 14, 22 ENVCUBE 18, 26
	float4 Color3;        // 15, 23 ENVCUBE 19, 27
	float4 Velocity;      // 16, 24 ENVCUBE 20, 28
	float4 Acceleration;  // 17, 25 ENVCUBE 21, 29
	float4 Wind;          // 18, 26 ENVCUBE 22, 30
}

float2x2 GetRotationMatrix(float angle)
{
	float sine, cosine;
	sincos(angle, sine, cosine);

	return float2x2(float2(cosine, -sine), float2(sine, cosine));
}

VS_OUTPUT main(VS_INPUT input)
{
	VS_OUTPUT vsout;

#	if defined(ENVCUBE)
#		if defined(RAIN)
	float2 positionOffset = input.TexCoord1.xy;
#		else
	float2x2 rotationMatrix = GetRotationMatrix(fVars0.w);
	float2 positionOffset = mul(rotationMatrix, input.TexCoord1.xy) * ScaleAdjust + mul(rotationMatrix, input.TexCoord1.zw);
#		endif

	float3 normalizedPosition = (fVars0.xyz + input.Position.xyz) / fVars2.xxx;
	normalizedPosition = normalizedPosition >= -normalizedPosition ? frac(abs(normalizedPosition)) :
	                                                                 -frac(abs(normalizedPosition));

	float4 msPosition;
	msPosition.xyz = normalizedPosition * fVars2.xxx + (-(fVars2.x * 0.5).xxx + fVars1.xyz);
	msPosition.w = 1;

	float4 viewPosition = mul(WorldViewProj, msPosition);
#		if defined(RAIN)
	float3 rainVelocity = Velocity.xyz;
#		if defined(EFFECTS11)
	if (SharedData::enbSettings.EnableRain) {
		float velLen = length(rainVelocity);
		if (velLen > 0) {
			float3 normVel = rainVelocity / velLen;
			rainVelocity = lerp(normVel, rainVelocity, SharedData::enbSettings.RainMotionStretch);
		}
	}
#		endif
	float4 adjustedMsPosition = msPosition - float4(rainVelocity, 0);
	float positionBlendParam = 0.5 * (1 + input.TexCoord1.y);
	float4 adjustedViewPosition = mul(WorldViewProj, adjustedMsPosition);
	float4 finalViewPosition = lerp(adjustedViewPosition, viewPosition, positionBlendParam);
#		else
	float4 finalViewPosition = viewPosition;
#		endif
	vsout.Position.xy = positionOffset + finalViewPosition.xy;
	vsout.Position.zw = finalViewPosition.zw;
	vsout.Color.xyz = 1.0.xxx;
	vsout.Color.w = fVars1.w;

	vsout.TexCoord0.xy = input.TexCoord0.xy;

	float4 precipitationOcclusionTexCoord = mul(PrecipitationOcclusionWorldViewProj, msPosition);
	precipitationOcclusionTexCoord.y = -precipitationOcclusionTexCoord.y;
	vsout.PrecipitationOcclusionTexCoord = precipitationOcclusionTexCoord;
#	else
	float tmp2 = input.Normal.w * input.Position.w;
	float tmp1 = tmp2 / fVars0.y;

	float uvScale1, uvScale2, tmp3, tmp4;
	if (tmp1 > fVars2.w) {
		uvScale1 = fVars2.y;
		uvScale2 = 0;
		tmp3 = fVars2.w;
		tmp4 = 1;
	} else if (tmp1 > fVars2.z) {
		uvScale1 = fVars2.x;
		uvScale2 = fVars2.y;
		tmp3 = fVars2.z;
		tmp4 = fVars2.w;
	} else {
		uvScale1 = 0;
		uvScale2 = fVars2.x;
		tmp3 = 0;
		tmp4 = fVars2.z;
	}
	float uvScaleParam = (tmp1 - tmp3) / (tmp4 - tmp3);
	float uvScale = lerp(uvScale1, uvScale2, uvScaleParam);

	vsout.TexCoord0.xy = fVars4.xy * input.TexCoord1.xy;

	float2 uv1 = (input.TexCoord1.zw * 2.0.xx - 1.0.xx) * uvScale;
	float uvAngle = input.TexCoord0.y * input.Position.w + input.TexCoord0.x;
	float2x2 rotationMatrix = GetRotationMatrix(uvAngle);
	float2 positionOffset = mul(rotationMatrix, uv1);

	float4 msPosition;
	msPosition.xyz = -fVars3.xyz +
	                 (((input.Normal.xyz * fVars0.www + Acceleration.xyz) * (tmp2 * tmp2)) * 0.5 +
						 (((fVars0.zzz * input.Normal.xyz) * input.TexCoord0.zzz +
							  (normalize(-Wind.xyz + input.Position.xyz) * Wind.www + Velocity.xyz)) *
								 tmp2 +
							 input.Position.xyz));
	msPosition.w = 1;

	float4 viewPosition = mul(WorldViewProj, msPosition);
	vsout.Position.xy = positionOffset * ScaleAdjust + viewPosition.xy;
	vsout.Position.zw = viewPosition.zw;

	float4 color1, color2;
	float colorTmp1, colorTmp2;
	if (tmp1 > fVars1.z) {
		color1 = Color3.xyzw;
		color2 = float4(Color3.xyz, 0);
		colorTmp1 = fVars1.z;
		colorTmp2 = 1;
	} else if (tmp1 > fVars1.y) {
		color1 = Color2.xyzw;
		color2 = Color3.xyzw;
		colorTmp1 = fVars1.y;
		colorTmp2 = fVars1.z;
	} else if (tmp1 > fVars1.x) {
		color1 = Color1.xyzw;
		color2 = Color2.xyzw;
		colorTmp1 = fVars1.x;
		colorTmp2 = fVars1.y;
	} else {
		color1 = float4(Color1.xyz, 0);
		color2 = Color1.xyzw;
		colorTmp1 = 0;
		colorTmp2 = fVars1.x;
	}
	float colorParam = (tmp1 - colorTmp1) / (colorTmp2 - colorTmp1);
	float4 color = lerp(color1, color2, colorParam);

	vsout.Color.w = fVars3.w * color.w;
	vsout.Color.xyz = color.xyz;
#	endif

#		if defined(ENVCUBE) && defined(RAIN) && defined(EFFECTS11)
	vsout.RaindropData.xy = input.TexCoord1.xy * 0.5 + 0.5;
#		endif

	return vsout;
}
#endif

typedef VS_OUTPUT PS_INPUT;

struct PS_OUTPUT
{
	float4 Color: SV_Target0;
	float4 Normal: SV_Target1;
};

#ifdef PSHADER

#	if defined(LIGHT_LIMIT_FIX)
#		include "LightLimitFix/LightLimitFix.hlsli"
#	endif

#	if defined(ISL) && defined(LIGHT_LIMIT_FIX)
#		include "InverseSquareLighting/InverseSquareLighting.hlsli"
#	endif

SamplerState SampSourceTexture : register(s0);
#	if defined(GRAYSCALE_TO_COLOR) || defined(GRAYSCALE_TO_ALPHA)
SamplerState SampGrayscaleTexture : register(s1);
#	endif
#	if defined(ENVCUBE)
SamplerState SampPrecipitationOcclusionTexture : register(s2);
SamplerState SampUnderwaterMask : register(s3);
#	endif

Texture2D<float4> TexSourceTexture : register(t0);
#	if defined(GRAYSCALE_TO_COLOR) || defined(GRAYSCALE_TO_ALPHA)
Texture2D<float4> TexGrayscaleTexture : register(t1);
#	endif
#	if defined(ENVCUBE)
Texture2D<float4> TexPrecipitationOcclusionTexture : register(t2);
Texture2D<float4> TexUnderwaterMask : register(t3);
#	endif
#	if defined(ENVCUBE) && defined(RAIN) && defined(EFFECTS11)
Texture2D<float4> TexRaindropNormals : register(t80);
#	endif

cbuffer PerGeometry : register(b2)
{
	float ColorScale : packoffset(c0);
	float3 TextureSize : packoffset(c1);
};

#	define LinearSampler SampSourceTexture
#	include "Common/ShadowSampling.hlsli"

#	if defined(IBL)
#		include "IBL/IBL.hlsli"
#	endif

#	if defined(DYNAMIC_CUBEMAPS)
#		define SampColorSampler SampSourceTexture
#		include "DynamicCubemaps/DynamicCubemaps.hlsli"
#	endif

PS_OUTPUT main(PS_INPUT input, bool frontFace : SV_IsFrontFace)
{
	PS_OUTPUT psout;

#	if defined(ENVCUBE)
	float2 precipitationOcclusionUV = (input.PrecipitationOcclusionTexCoord.xy * 0.5 + 0.5) * TextureSize.x;
	float precipitationOcclusion = -input.PrecipitationOcclusionTexCoord.z + TexPrecipitationOcclusionTexture.Load(float3(precipitationOcclusionUV, 0)).x;
	float2 underwaterMaskUv = TextureSize.yz * input.Position.xy;
	float underwaterMask = TexUnderwaterMask.Sample(SampUnderwaterMask, underwaterMaskUv).x;
	if (precipitationOcclusion - underwaterMask < 0) {
		discard;
	}
#	endif

#	if defined(ENVCUBE) && defined(RAIN) && defined(DYNAMIC_CUBEMAPS) && defined(EFFECTS11)
if (SharedData::enbSettings.EnableRain) {
	float4 raindropNormal = TexRaindropNormals.Sample(SampSourceTexture, input.RaindropData.xy);
    float alpha = saturate(raindropNormal.w * (1.0 - SharedData::enbSettings.RainMotionTransparency));
   	clip(alpha - (4.0 / 255.0));
	raindropNormal.y = 1.0 - raindropNormal.y;

    // Reconstruct camera-relative worldspace position (camera at origin).
    float2 uv = input.Position.xy * SharedData::BufferDim.zw;
    float4 posCS = float4(2.0 * float2(uv.x, 1.0 - uv.y) - 1.0, input.Position.z, 1.0);
    float4 posWS = mul(FrameBuffer::CameraViewProjInverse, posCS);
    posWS.xyz /= posWS.w;

    // Build worldspace TBN from screen-space derivatives. The billboard is camera-aligned,
    // so dPdx/dPdy lie in the billboard plane along screen X/Y.
    float3 T = normalize(ddx(posWS.xyz));
    float3 N = normalize(cross(T, -ddy(posWS.xyz)));
    float3 B = cross(N, T);
    float3x3 TBN = float3x3(T, B, N);

    float3 normalTS = normalize(raindropNormal.xyz * 2.0 - 1.0);
    float3 normalWS = normalize(mul(normalTS, TBN));

	if (frontFace)
		normalWS = -normalWS;

    float3 V = normalize(-posWS.xyz);
    float NdotV = saturate(dot(normalWS, V));
    float fresnel = 0.02 + 0.98 * pow(1.0 - NdotV, 5.0);

    float3 reflectDir = reflect(-V, normalWS);
    float3 refractDir = refract(-V, normalWS, 1.0 / 1.33);
    if (dot(refractDir, refractDir) < 1e-4)
        refractDir = -V;

    float3 reflectColor = Color::IrradianceToLinear(DynamicCubemaps::EnvReflectionsTexture.SampleLevel(SampSourceTexture, reflectDir, 0).xyz);
    float3 refractColor = Color::IrradianceToLinear(DynamicCubemaps::EnvReflectionsTexture.SampleLevel(SampSourceTexture, refractDir, 0).xyz);

    psout.Color.xyz = Color::IrradianceToGamma(lerp(refractColor, reflectColor, fresnel));
    psout.Color.w = alpha;
    psout.Normal = float4(0, 1, 0, alpha);
    return psout;
}
#	endif

	float4 sourceColor = TexSourceTexture.Sample(SampSourceTexture, input.TexCoord0);
	float4 baseColor = input.Color * sourceColor;
	baseColor.xyz = Color::Diffuse(baseColor.xyz);
#	if defined(GRAYSCALE_TO_COLOR)
	float3 grayScaleColor =
		TexGrayscaleTexture.Sample(SampGrayscaleTexture, float2(sourceColor.y, input.Color.x)).xyz;
	baseColor.xyz = grayScaleColor;
#	endif
#	if defined(GRAYSCALE_TO_ALPHA)
	float grayScaleAlpha =
		TexGrayscaleTexture.Sample(SampGrayscaleTexture, float2(sourceColor.w, input.Color.w)).w;
	baseColor.w = grayScaleAlpha;
#	endif

	float3 propertyColor = 0.0;

	float2 uv = input.Position.xy * SharedData::BufferDim.zw;

	float4 positionWS = float4(2 * float2(uv.x, -uv.y + 1) - 1, input.Position.z, 1);
	positionWS = mul(FrameBuffer::CameraViewProjInverse, positionWS);
	positionWS.xyz = positionWS.xyz / positionWS.w;

	float unusedDetailedShadow;
	float3 dirLightColor = SharedData::DirLightColor.xyz * ShadowSampling::GetLightingShadow(positionWS.xyz, unusedDetailedShadow);
	float3 ambientColor = max(0, SharedData::GetAmbient(float3(0, 0, 1)));
#	if defined(IBL)
	if (SharedData::iblSettings.EnableIBL) {
		ambientColor = ImageBasedLighting::GetDiffuseIBL(ambientColor, float3(0, 0, -1));
	}
#	endif

	propertyColor += dirLightColor;
	propertyColor += ambientColor;

#	if defined(LIGHT_LIMIT_FIX)
	uint lightCount = 0;
	{
		float3 viewPosition = FrameBuffer::WorldToView(positionWS.xyz);
		float2 screenUV = FrameBuffer::ViewToUV(viewPosition);

		uint clusterIndex = 0;
		if (LightLimitFix::GetClusterIndex(screenUV, viewPosition.z, clusterIndex)) {
			lightCount = LightLimitFix::lightGrid[clusterIndex].lightCount;
			uint lightOffset = LightLimitFix::lightGrid[clusterIndex].offset;
			[loop] for (uint i = 0; i < lightCount; i++)
			{
				uint clusteredLightIndex = LightLimitFix::lightList[lightOffset + i];
				LightLimitFix::Light light = LightLimitFix::lights[clusteredLightIndex];
				if (LightLimitFix::IsLightIgnored(light) || light.lightFlags & LightLimitFix::LightFlags::Shadow) {
					continue;
				}
				float3 lightDirection = light.positionWS.xyz - positionWS.xyz;
				float lightDist = length(lightDirection);

#		if defined(ISL)
				float intensityMultiplier = InverseSquareLighting::GetAttenuation(lightDist, light);
#		else
				float intensityFactor = saturate(lightDist / light.radius);
				float intensityMultiplier = 1 - intensityFactor * intensityFactor;
#		endif

				float3 lightColor = light.color.xyz * intensityMultiplier;
				propertyColor += lightColor;
			}
		}
	}
#	endif

	psout.Color.xyz = propertyColor * baseColor.xyz;
	psout.Color.w = baseColor.w;
	psout.Normal.w = baseColor.w;
	psout.Normal.xyz = float3(0, 1, 0);

	return psout;
}
#endif
