Texture2D<float4> Texture0 : register(t0, space2);
SamplerState Sampler0 : register(s0, space2);

Texture2D<float4> Texture1 : register(t1, space2);
SamplerState Sampler1 : register(s1, space2);

Texture2D<float4> Texture2 : register(t2, space2);
SamplerState Sampler2 : register(s2, space2);

Texture2D<float> DirectionalShadowDepthMap : register(t3, space2);
SamplerState DirectionalShadowSampler : register(s3, space2);

cbuffer Uniforms : register(b0, space3)
{
    float4 viewPos;
    float4 lightPosition;
    float4 directionalLight;
    float4 directionalLightColor;
    float2 screenSize;
    float time;
};

struct Scene
{
    float4x4 Model;

    uint textureMode;    // 0 = regular triple, 
                        // 0x10 = video yuv,
    uint pad0;    // 0 = regular triple, 
    uint pad1;    // 0 = regular triple, 
    uint pad2;    // 0 = regular triple, 
};

StructuredBuffer<Scene> scene: register(t4, space2);

// from learnopengl.com
float3 BlinnPhong(float3 normal, float3 fragPos, float3 lightPos, float3 lightColor)
{
    float dist = length(lightPos - fragPos);
    // attenuation
    float maxDist = 15;
    float cutoff1 = 2;
    float cutoff2 = 4;
    float cutoffFactor = 0.5;
    float attenuation = 1.0 / (dist * dist);

    if (dist > cutoff1)
    {
        attenuation *= cutoffFactor;
    }
    if (dist > cutoff2)
    {
        attenuation *= cutoffFactor;
    }
    if (dist > maxDist)
    {
        attenuation = 0;
    }

    if (attenuation > 1.0)
    {
        attenuation = 1.0;
    }

    if(length(normal) < 0.1)
    {
        return fragPos * lightColor * attenuation;
    }

    // diffuse  parameter
    float3 lightDir = normalize(lightPos - fragPos);
    float diff = max(dot(lightDir, normal), 0.0);
    float3 diffuse = diff * lightColor;

    // specular parameter
    float3 viewDir = normalize(fragPos - viewPos.xyz);
    float3 reflectDir = reflect(-lightDir, normal);
    float spec = 0.0;
    float3 halfwayDir = normalize(lightDir + viewDir);
    spec = pow(max(dot(normal, halfwayDir), 0.0), 30.0);
    float3 specular = spec * lightColor * 2;

    

    diffuse *= attenuation;
    specular *= attenuation;

    return diffuse + specular;
    
    //return lightColor * 100 * attenuation;
    
}

float shadowCalculation(float3 fragPos, float4 DirectionalShadowFragPos)
{
    float3 shadowProjection = DirectionalShadowFragPos.xyz / DirectionalShadowFragPos.w;
    shadowProjection = shadowProjection * 0.5 + 0.5 ;
    shadowProjection.x = shadowProjection.x;

    float2 uv = shadowProjection.xy;
    uv.y = -uv.y;
    float depthSample = DirectionalShadowDepthMap.Sample(DirectionalShadowSampler, uv);

    float fragDepth = DirectionalShadowFragPos.z;

    float bias = 0.005;
    float shadow = fragDepth - bias > depthSample ? 1.0 : 0.0;

    return shadow;
}

float shadowCalculationPCF(float3 fragPos, float4 DirectionalShadowFragPos)
{
    // size = 2048 x 2028
    float2 texDim = float2(2048, 2048);

	float scale = 0.5;
	float dx = scale * 1.0 / float(texDim.x);
	float dy = scale * 1.0 / float(texDim.y);

	float combinedDepth = 0.0;
	int count = 0;
	int range = 4;

    float3 shadowProjection = DirectionalShadowFragPos.xyz / DirectionalShadowFragPos.w;
    shadowProjection = shadowProjection * 0.5 + 0.5 ;
    shadowProjection.x = shadowProjection.x;

    float2 uv = shadowProjection.xy;
    uv.y = -uv.y;
    float fragDepth = DirectionalShadowFragPos.z;
    float combinedShadow = 0.0;
    float bias = 0.005;

    for (int y = -range ; y <= range ; y++) {
        for (int x = -range ; x <= range ; x++) {
            float2 Offsets = float2(x * dx, y * dy);
            //float3 UVC = float3(uv + Offsets, DirectionalShadowFragPos.z);
            float depth = DirectionalShadowDepthMap.Sample(DirectionalShadowSampler, uv + Offsets);
            combinedDepth += depth;
            combinedShadow += fragDepth - bias > depth ? 1.0 : 0.0;
            count += 1;
        }
    }

    combinedDepth = combinedDepth / (count);
    combinedShadow = combinedShadow / (count);

    // float fragDepth = DirectionalShadowFragPos.z;
    // float shadow = fragDepth - bias > combinedDepth ? 1.0 : 0.0;

    return combinedShadow;
}

float3 DirectionalLight(float3 normal, float3 fragPos, float4  DirectionalShadowFragPos)
{

    float3 lightDir = directionalLight.xyz;
    float diff = max(dot(lightDir, normal), 0.0);
    float3 diffuse = diff * directionalLightColor.xyz;

    // specular parameter
    float3 viewDir = normalize(fragPos - viewPos.xyz);
    float3 reflectDir = reflect(-lightDir, normal);
    float spec = 0.0;
    float3 halfwayDir = normalize(lightDir + viewDir);
    spec = pow(max(dot(normal, halfwayDir), 0.0), 30.0);
    float3 specular = spec * directionalLightColor.xyz * 2;

    // shadow mapping
    
    diffuse = diffuse * (1.0 - shadowCalculationPCF(fragPos, DirectionalShadowFragPos));

    return diffuse;
}



float4 main(
    float4 ScreenPosition: SV_POSITION,

    float2 UV : TEXCOORD0,
    float3 WorldPos: TEXCOORD1,
    float3 Normal: TEXCOORD2,
    float4 DirectionalShadowFragPos: TEXCOORD3,
    uint Instance: TEXCOORD4
) : SV_Target0
{


    float4 s; 

    int texMode = scene[Instance].textureMode;
    if ((texMode & 0xf0) == 0x10) // YUV
    {
        float3 offset = float3(-0.0625, -0.5, -0.5);
        float3 Rcoeff = float3(1.164,  0.000,  1.793);
        float3 Gcoeff = float3(1.164, -0.213, -0.533);
        float3 Bcoeff = float3(1.164,  2.112,  0.000);

        float3 yuv; 
        float3 rgb;
        float2 uv = UV;

        yuv.x = Texture0.Sample(Sampler0, uv).r;
        yuv.y = Texture1.Sample(Sampler1, uv).r;
        yuv.z = Texture2.Sample(Sampler2, uv).r;
        yuv += offset;

        rgb.x = dot(yuv, Rcoeff);
        rgb.y = dot(yuv, Gcoeff);
        rgb.z = dot(yuv, Bcoeff);

        s = float4(rgb, 1.0); 
    }
    else if((texMode & 0xf0) == 0x0)
    {
        s = Texture0.Sample(Sampler0, UV);
    }

    if((texMode & 0x1) == 1) // srgb correction bit
    {
        s = pow(s, 2.2);
    }

    float alpha = s.w;
    if(alpha < 0.01)
    {
        discard;
    }

    float3 lightPos = lightPosition.xyz; //+ float3(0, 20, 0);
    float3 lightColor = float3(0.8, 0.8, 0.5);

    float3 ambient = float3(0.1,0.1,0.1);


    float3 col = s.xyz;
    float3 c2 = col * ambient 
        + col * BlinnPhong(Normal, WorldPos, lightPos, lightColor) 
        + col * DirectionalLight(Normal, WorldPos,  DirectionalShadowFragPos);


    float4 rv = float4(pow(c2, 1.0 / 2.2), alpha);

// === debug bone zone
    // float3 shadowProjection = DirectionalShadowFragPos.xyz / DirectionalShadowFragPos.w;
    // shadowProjection = shadowProjection * 0.5 + 0.5 ;
    // shadowProjection.x = shadowProjection.x;

    // float2 uv = shadowProjection.xy;
    // uv.y = -uv.y;
    // float depthSample = DirectionalShadowDepthMap.Sample(DirectionalShadowSampler, uv);
    // float fragDepth = shadowProjection.z;

    // float t = shadowCalculationPCF(WorldPos, DirectionalShadowFragPos);


    // float t = depthSample;
    // rv = float4(t, t, t, 1.0);

    // rv.x = rv.x + lightPos.x * 0.0001;

    return rv;
}