﻿// Copyright (c) .NET Foundation and Contributors (https://dotnetfoundation.org/ & https://stride3d.net) and Silicon Studio Corp. (https://www.siliconstudio.co.jp)
// Distributed under the MIT license. See the LICENSE.md file in the project root for more information.
namespace Stride.Rendering.Materials
{
    /// <summary>
    /// Performs hair shading.
    /// </summary>
    shader MaterialSurfaceShadingSpecularHair<int ShadingModel, bool DebugRenderPasses> :
        IMaterialSurfaceShading, // Required for "ComputeDirectLightContribution()" and the like.
        MaterialPixelShadingStream,
        MaterialHairShared,
        NormalStream, // Required for "streams.normalWS", "streams.meshNormal", "streams.meshTangent".
        Transformation, // Required for "WorldInverseTranspose".
        MaterialPixelStream    // Required for "streams.matDiffuse".
    {
        compose ComputeColor SpecularHighlightsShiftNoiseTexture;
        compose ComputeColor SecondarySpecularGlintsNoiseTexture;
        
        compose IMaterialHairLightAttenuationFunction hairLightAttenuationFunction;
        compose IMaterialHairDirectionFunction hairDirectionFunction;
        compose IMaterialHairShadowingFunction hairShadowingFunction;
        compose IMaterialHairDiscardFunction hairDiscardFunction;
        compose IMaterialSpecularMicrofacetEnvironmentFunction environmentFunction;
        
        cbuffer PerMaterial
        {
            [Color]
            stage float3 HairSpecularColor1;    // The color of the primary specular reflection.
      
            [Color]
            stage float3 HairSpecularColor2;    // The color of the secondary specular reflection.
            
            stage float HairScalesAngle;
            stage float HairSpecularShiftRatio;
            stage float HairSpecularExponent1;
            stage float HairSpecularExponent2;
            stage float HairSpecularScale1;
            stage float HairSpecularScale2;
            stage float HairShiftNoiseScale;    // Controls how much the noise should affect the specular highlight direction.
            stage float HairGlintsNoiseStrength;    // Controls how much the glints noise should affect the secondary reflections.
        }
        
        /// <summary>
        /// Holds all the required input data for hair shading.
        /// </summary>
        struct SurfaceData
        {
            float hairScalesAngle;
            float hairSpecularShiftRatio;
            float3 WBinormal;
            float3 WViewDir;
            float3 WNormal;
            float hairSpecularExponent1;    // TODO: This is ignored by the indirect lighting.
            float hairSpecularExponent2;    // TODO: This is ignored by the indirect lighting.
            float3 hairSpecularColor1;
            float3 hairSpecularColor2;
            float cavity;
            float hairSpecularScale1;
            float hairSpecularScale2;
            float hairSecondarySpecularGlintsNoise;
            float hairSpecularHighlightsShiftNoise;
        };

        //stream SurfaceData surfaceData;
        static SurfaceData surfaceData;    // TODO: How to make the variable not get exposed in the CodeBehind? Sadly making this a streams variable causes a shader compilation error.

        SurfaceData GenerateSurfaceData(void)
        {
            SurfaceData result;
            result.hairScalesAngle = HairScalesAngle;
            result.hairSpecularShiftRatio = HairSpecularShiftRatio;
            result.WViewDir = streams.viewWS;
            result.WNormal = streams.normalWS;
            
            // TODO: Use mesh normal or normal map normal?
            result.WBinormal = hairDirectionFunction.Compute(streams.meshNormal, streams.meshTangent, (float3x3)WorldInverseTranspose); // TODO: PERFORMANCE: Get rid of the float3x3 cast?
            
            result.hairSpecularExponent1 = HairSpecularExponent1;
            result.hairSpecularExponent2 = HairSpecularExponent2;

            result.hairSpecularColor1 = HairSpecularColor1;
            result.hairSpecularColor2 = HairSpecularColor2;

            result.hairSpecularScale1 = HairSpecularScale1;
            result.hairSpecularScale2 = HairSpecularScale2;

            // We can't use "streams.matSpecularVisible" here because it is multiplied by the specular color.
            result.cavity = lerp(1.0, streams.matCavity, streams.matCavitySpecular);
            
            result.hairSpecularHighlightsShiftNoise = (SpecularHighlightsShiftNoiseTexture.Compute().r - 0.5) * HairShiftNoiseScale;
            //result.hairSpecularHighlightsShiftNoise = SpecularHighlightsShiftNoiseTexture.Compute().r * HairShiftNoiseScale;
            result.hairSecondarySpecularGlintsNoise = lerp(1.0, SecondarySpecularGlintsNoiseTexture.Compute().r, HairGlintsNoiseStrength);
            
            return result;
        }

        override void PrepareForLightingAndShading()    // This gets executed only once for the entire shader.
        {
            surfaceData = GenerateSurfaceData();
        }

        void CalculateShiftAngles(SurfaceData surfaceData, out float shiftAngle1, out float shiftAngle2)
        {
            // The hair shift is being calculated differently from Mizuchi because the Mizuchi implementation is weird.

            float scalesAngle = surfaceData.hairScalesAngle;
            shiftAngle1 = 2.0 * scalesAngle;
            shiftAngle2 = -shiftAngle1 * surfaceData.hairSpecularShiftRatio;  // hairSpecularShiftRatio is theoretically 1.5

            shiftAngle1 += surfaceData.hairSpecularHighlightsShiftNoise;
            shiftAngle2 += surfaceData.hairSpecularHighlightsShiftNoise;    // I think Mizuchi subtracts the noise from the 2nd shift angle. That isn't correct according to my research, so I changed it.
            
        }

        float HairSingleSpecularTerm_Kajiya(float3 T, float3 toLightVec, float3 viewDir, float shiftAngle)
        {
            float cosTL = dot(T, toLightVec);
            float sinTL = sqrt(1.0 - cosTL * cosTL);

            // in Kajiya's model: specular component: cos(t, rl) * cos(t, e) + sin(t, rl)sin(t, e)
            float cosT_RL = -cosTL;
            float sinT_RL = sinTL;
            float cosTE   = dot(T, viewDir);
            float sinTE   = sqrt(1.0 - cosTE * cosTE);

            // Kajiya-Kay highlight: reflected direction shifted
            float cosT_RL_shifted = cosT_RL * cos(shiftAngle) - sinT_RL * sin(shiftAngle);
            float sinT_RL_shifted = sqrt(1 - cosT_RL_shifted * cosT_RL_shifted);
            return max(0.0, cosT_RL_shifted * cosTE + sinT_RL_shifted * sinTE);
        }

        //------------------------------------------------------------------------------
        // From Shader X3, Chapter 2.14 (Hair Rendering and Shading), P.244 :
        // "sum the contribution of two separate specular terms per light"
        // "each term has different specular colors, exponents, and is shifted in different directions
        //------------------------------------------------------------------------------
        float HairSingleSpecularTerm_Scheuermann(float3 T, float3 H, float exponent)
        {
            float dotTH = dot(T, H);
            //float sinTH = sqrt(1.0 - dotTH * dotTH);  // Original version from Mizuchi. Causes artifacts with Scheuermann approximation because of the mathematically correct rotation. The original implementation doesn't suffer from that issue.
            float sinTH = sqrt(max(1.0 - dotTH * dotTH, 0.0));  // We limit the value above 0.0 so it doesn't cause NaN errors with the Scheuermann approximation (because of the mathematically correct rotation).
            return pow(sinTH, exponent);
        }

        //==============================================================================
        // From Shader X3, Chapter 2.14 (Hair Rendering and Shading), P.244 :
        // "specular highlights are shifted because the scales on the surface of a hair strand have tilted normals"
        // "ideally we would tilt the tangent in the direction of the viewer"
        // "in practice it is sufficient to tilt in the direction of the geometric normal"
        //==============================================================================
        //------------------------------------------------------------------------------
        // Original formula
        //------------------------------------------------------------------------------
        //float3 ShiftTangent(float3 T, float3 N, float shiftAmount)
        //{
        //    return normalize(T + shiftAmount * N);
        //}
        //
        //------------------------------------------------------------------------------
        // Alternative formula (mathematically correct rotation):
        // Return vector X (= vector T rotated by angle shiftAngle towards direction N)
        //------------------------------------------------------------------------------
        float3 ShiftTangent(float3 T, float3 N, float shiftAngle)
        {
            // While this function performs mathematically correct rotation, it's probably not desired
            // (because the shift texture doesn't represent angles) and instead the original formula should be preferred.
            // This is because it can cause the rotation to go beyond 90 degrees, causing the shading vector to point inside the surface,
            // causing weird artifacts and issues in other parts of the shading code, which weren't modified to work with the mathematically
            // correct rotation but still assume the original shift implemnetation.

            float cosTX = cos(shiftAngle);
            float sinTX = sqrt(1.0 - cosTX * cosTX) * sign(shiftAngle);

            if(ShadingModel == HAIR_SHADING_SCHEUERMANN_APPROXIMATION)
            {
                //return normalize(N + T * shiftAngle);   // Original formula
                // Simplification (if T,N normal) - Use when T and N are normal
                //float cosTN = 0.0;
                //float sinTN = 1.0;
                return cosTX * T + sinTX * N;
            }
            else
            {
                // Handling case when T and N are not really normal
                float cosTN = dot(T, N);
                float sinTN = sqrt(1.0 - cosTN * cosTN);
                float3 X = (cosTX * sinTN - cosTN * sinTX) * T + sinTN * sinTX * N;
                return normalize(X);
            }
        }

        float HairSpecularAttenuation(float dotNL)  // "dotNL" must be unsaturated!
        {
            if(ShadingModel == HAIR_SHADING_KAJIYAKAY_SHIFTED)
            {
                // No ad-hoc attenuation in Kajiya-Kay formula
                // The coefficient for diffuse is applied during ComputeHairLightingSpecular(),
                // therefore we do not apply dotNL.
                return 1.0;
            }
            else
            {
                //------------------------------------------------------------------------------
                // From Shader X3, Chapter 2.14 (Hair Rendering and Shading), P.246 :
                // "specular attenuation for hair facing away from light"
                //------------------------------------------------------------------------------
                return saturate(1.75 * dotNL + 0.25);
            }
        }

        float3 ComputeSpecularKajiya(SurfaceData surfaceData, float3 toLightVec, float shiftAngle,
                                     float specularScale, float specularExponent, float3 specularColor)
        {
            //--------------------------------------------------------------------------
            // Shifted Kajiya-Kay formula (extension of Kajiya-Kay)
            // Similar to what is used in AMD TressFX sample.
            //--------------------------------------------------------------------------

            float shiftedSpecular = HairSingleSpecularTerm_Kajiya(surfaceData.WBinormal, toLightVec, surfaceData.WViewDir, shiftAngle);
            return(specularScale * specularColor * pow(shiftedSpecular, specularExponent));
        }
        
        float3 ComputeSpecularScheuermann(SurfaceData surfaceData, float3 halfVector, float shiftAngle,
                                          float specularScale, float specularExponent, float3 specularColor)
        {
            //--------------------------------------------------------------------------
            // Scheuermann formula
            //--------------------------------------------------------------------------
            // From Shader X3, Chapter 2.14 (Hair Rendering and Shading), P.245 :
            // "The lack of real self-shadowing will cause the specular highlights to be
            //  too bright on the hair facing away from the light"
            // "To fade out the specular highlight on the shadowed side of the hair, we 
            //  multiply by an attenuation term that is similar to the diffuse term."
            //--------------------------------------------------------------------------

            // shift tangents
            float3 shiftDirection;

            if(ShadingModel == HAIR_SHADING_SCHEUERMANN_APPROXIMATION)
            {
                shiftDirection = surfaceData.WNormal;
            }
            else if(ShadingModel == HAIR_SHADING_SCHEUERMANN_IMPROVED)
            {
                shiftDirection = surfaceData.WViewDir;
            }

            const float3 T = ShiftTangent(surfaceData.WBinormal, shiftDirection, shiftAngle);

            // specular term
            return(specularScale * specularColor * HairSingleSpecularTerm_Scheuermann(T, halfVector, specularExponent));
        }
                
        //==============================================================================
        // Specular term for hair (for Direct Lighting)
        //==============================================================================
        float3 ComputeHairLightingSpecular(SurfaceData surfaceData, float3 halfVector, float3 toLightVec)
        {
            // modulate specular shift by a texture
            float shiftAngle1;
            float shiftAngle2;
            CalculateShiftAngles(surfaceData, shiftAngle1, shiftAngle2);
            
            float hairSpecularExponent1 = surfaceData.hairSpecularExponent1;    // We already make sure on the CPU side that "surfaceData.hairSpecularExponent1" can't be "0.0".
            float hairSpecularExponent2 = surfaceData.hairSpecularExponent2;    // We already make sure on the CPU side that "surfaceData.hairSpecularExponent2" can't be "0.0".
            
            float3 specular1;
            float3 specular2;
            if(ShadingModel == HAIR_SHADING_KAJIYAKAY_SHIFTED)
            {
                // primary highlight: reflected direction shift towards tip
                specular1 = ComputeSpecularKajiya(surfaceData, toLightVec, shiftAngle1, surfaceData.hairSpecularScale1, hairSpecularExponent1, surfaceData.hairSpecularColor1);
                // secondary highlight: reflected direction shifted toward root
                specular2 = ComputeSpecularKajiya(surfaceData, toLightVec, shiftAngle2, surfaceData.hairSpecularScale2, hairSpecularExponent2, surfaceData.hairSpecularColor2);
            }
            else
            {
                // specular terms
                specular1 = ComputeSpecularScheuermann(surfaceData, halfVector, shiftAngle1, surfaceData.hairSpecularScale1, hairSpecularExponent1, surfaceData.hairSpecularColor1);
                specular2 = ComputeSpecularScheuermann(surfaceData, halfVector, shiftAngle2, surfaceData.hairSpecularScale2, hairSpecularExponent2, surfaceData.hairSpecularColor2);
            }
            
            // modulate secondary specular term with noise
            specular2 *= surfaceData.hairSecondarySpecularGlintsNoise;
            
            float3 specular = (specular1 + specular2) * hairShadowingFunction.Compute(); // Used to apply shadowing/scattering.
            return specular * surfaceData.cavity;
        }

        float3 ComputeSpecularIndirectLighting(SurfaceData surfaceData, float shiftAngle, float3 specularColor)
        {
            const float3 shiftedNormal = ShiftTangent(surfaceData.WBinormal, surfaceData.WNormal, shiftAngle);
            
            // specular term
            const float shiftedNdotV = max(dot(shiftedNormal, surfaceData.WViewDir), 0.0);
            
            float3 specular = environmentFunction.Compute(specularColor, streams.alphaRoughness, shiftedNdotV) * streams.envLightSpecularColor;
            
            // Since "SpecularTerm_IBL()", which is replaced by "environmentFunction.Compute()",
            // applies the cavity parameter internally, we do it too:
            specular *= surfaceData.cavity;
            
            return specular;
        }
       
        //==============================================================================
        // Specular term for hair (for Indirect Lighting)
        //------------------------------------------------------------------------------
        // Similarly to what is done for the direct lighting, we compute the specular component based on the shifted hair tangents.
        // This is not physically realistic because:
        // - the ambientBRDF map does not convolve the hair specular BRDF
        // However, this is a satisfying approximation at first.
        //==============================================================================
        float3 ComputeHairImageBasedLightingSpecular(SurfaceData surfaceData)   //, IBLTextureParameters IBL)
        {
            //return 1.0;

            // shift tangents
            float shiftAngle1;
            float shiftAngle2;
            CalculateShiftAngles(surfaceData, shiftAngle1, shiftAngle2);
            
            float3 specular1 = ComputeSpecularIndirectLighting(surfaceData, shiftAngle1, surfaceData.hairSpecularColor1);
            float3 specular2 = ComputeSpecularIndirectLighting(surfaceData, shiftAngle2, surfaceData.hairSpecularColor2);
            
            specular1 *= surfaceData.hairSpecularScale1;    // TODO: This is a workaround to make the indirect lighting look better. Not sure if we should keep it as it basically scales the specular reflection like it does for the direct lighting.
            specular2 *= surfaceData.hairSpecularScale2;    // TODO: This is a workaround to make the indirect lighting look better. Not sure if we should keep it as it basically scales the specular reflection like it does for the direct lighting.
            
            // modulate secondary specular term with noise
            specular2 *= surfaceData.hairSecondarySpecularGlintsNoise;

            return specular1 + specular2;
        }
        
        override float3 ComputeDirectLightContribution()
        {
            if(DebugRenderPasses)
            {
                return 0.0;   // Return 0.0 because the indirect lighting function already returns the debug color.
            }

            float3 specular = ComputeHairLightingSpecular(surfaceData,
                                                          streams.H,    // Half vector
                                                          streams.lightDirectionWS);    // Vector pointing from the pixel towards the light.
            
            const float NdotLUnsaturated = dot(normalize(streams.normalWS), normalize(streams.lightDirectionWS));   // TODO: PERFORMANCE: Normalization necessary?
            
            //specular *= saturate(NdotLUnsaturated * 0.5 + 0.5);   // This could be used in combination with hair SSS.
            specular *= HairSpecularAttenuation(NdotLUnsaturated);
            specular *= hairLightAttenuationFunction.Compute();
            specular *= streams.lightColor * streams.lightAttenuation * streams.matDiffuseSpecularAlphaBlend.y;
            //specular *= streams.lightDirectAmbientOcclusion;
            
            // TODO: Multiply by alpha or not?
            return specular * streams.matDiffuse.a; // TODO: Technically we should use "streams.shadingColorAlpha" but its value is assigned AFTER the shading, which makes it useless.
        }

        override float3 ComputeEnvironmentLightContribution()
        {
            if(DebugRenderPasses)
            {
                return GetDebugColor(PassID);
            }

            // TODO: Multiply by alpha or not?
            return ComputeHairImageBasedLightingSpecular(surfaceData).rgb * streams.matDiffuse.a;   // TODO: Technically we should use "streams.shadingColorAlpha" but its value is assigned AFTER the shading, which makes it useless.
        }

        override void AfterLightingAndShading()
        {
            hairDiscardFunction.Discard();
        }
    };
}