//Unity Toon Shader/HDRP
//nobuyuki@unity3d.com
//toshiyuki@unity3d.com (Universal RP/HDRP) 

#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Macros.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/PhysicalCamera.hlsl"
#include "Packages/com.misaki.hdrp-toon/Runtime/Shaders/Includes/Common/UtsMaterialEvaluation.hlsl"
#include "Packages/com.misaki.hdrp-toon/Runtime/Shaders/Includes/Lighting/UtsLightEvaluation.hlsl"
#include "Packages/com.misaki.hdrp-toon/Runtime/Models/SurfaceFeatureFlags.cs.hlsl"

// Channel mask enum.
// this must be same to UI cs code
// HDRPToonGUI._ChannelEnum
int eBaseColor = 0;
int eFirstShade = 1;
int eSecondShade = 2;
int eHighlight = 3;
int eAngelRing = 4;
int eRimLight = 5;
int eOutline = 6;

int GetNextDirectionalLightIndex(BuiltinData builtinData, int currentIndex, int mainLightIndex)
{
    int i = 0; // Declare once to avoid the D3D11 compiler warning.
    for (i = 0; i < (int)_DirectionalLightCount; ++i)
    {
        if (IsMatchingLightLayer(_DirectionalLightDatas[i].lightLayers, builtinData.renderingLayers))
        {
            if (mainLightIndex != i)
            {
                if (currentIndex < i)
                {
                    return i;
                }
            }
        }
    }
    return -1; // not found
}

int GetUtsMainLightIndex(BuiltinData builtinData)
{
    int mainLightIndex = -1;
    float3 lightColor = float3(0.0f, 0.0f, 0.0f);
    float  lightAttenuation = 0.0f;
    uint i = 0; // Declare once to avoid the D3D11 compiler warning.
    for (i = 0; i < _DirectionalLightCount; ++i)
    {
        if (IsMatchingLightLayer(_DirectionalLightDatas[i].lightLayers, builtinData.renderingLayers))
        {
            float3 currentLightColor = _DirectionalLightDatas[i].color;
            float  currentLightAttenuation = GetColorAttenuation(currentLightColor);

            if (mainLightIndex == -1 || (currentLightAttenuation > lightAttenuation))
            {
                mainLightIndex = i;
                lightAttenuation = currentLightAttenuation;

                lightColor = currentLightColor;
            } 
        }
    }

    return mainLightIndex;
}

bool UtsUseScreenSpaceShadow(DirectionalLightData light, float3 normalWS)
{
    #if defined(RAY_TRACED_SCREEN_SPACE_SHADOW_FLAG)
    // Two different options are possible here
    // - We have a ray trace shadow in which case we have no valid signal for a transmission and we need to fallback on the rasterized shadow
    // - We have a screen space shadow and it already contains the transmission shadow and we can use it straight away
    bool visibleLight = 0.5 * dot(normalWS, -light.forward) + 0.5 > 0.0;
    bool validScreenSpaceShadow = (light.screenSpaceShadowIndex & SCREEN_SPACE_SHADOW_INDEX_MASK) != INVALID_SCREEN_SPACE_SHADOW;
    bool rayTracedShadow = (light.screenSpaceShadowIndex & RAY_TRACED_SCREEN_SPACE_SHADOW_FLAG) != 0.0;
    return (validScreenSpaceShadow && ((rayTracedShadow && visibleLight) || !rayTracedShadow));
    #else
    return ( (light.screenSpaceShadowIndex & SCREEN_SPACE_SHADOW_INDEX_MASK) != INVALID_SCREEN_SPACE_SHADOW);
    #endif    
}

void UtsLightLoop(FragInputs fragInputs, PositionInputs posInput, UtsBSDFData bsdfData, BuiltinData builtinData,
    float3 V, uint featureFlags, out LightLoopOutput lightLoopOutput)
{
    LightLoopContext context;
    context.shadowContext = InitShadowContext();
    context.shadowValue = 1;
    context.sampleReflection = 0;
    #ifdef APPLY_FOG_ON_SKY_REFLECTIONS
        context.positionWS = posInput.positionWS;
    #endif

    // Initialize the contactShadow and contactShadowFade fields
    InitContactShadow(posInput, context);

    // First of all we compute the shadow value of the directional light to reduce the VGPR pressure
    if (featureFlags & LIGHTFEATUREFLAGS_DIRECTIONAL)
    {
        // Evaluate sun shadows.
        if (_DirectionalShadowIndex >= 0)
        {
            DirectionalLightData light = _DirectionalLightDatas[_DirectionalShadowIndex];

        #if defined(SCREEN_SPACE_SHADOWS_ON) && !defined(_SURFACE_TYPE_TRANSPARENT)
            if (UseScreenSpaceShadow(light, bsdfData.normalWS))
            {
                context.shadowValue = GetScreenSpaceColorShadow(posInput, light.screenSpaceShadowIndex).SHADOW_TYPE_SWIZZLE;
            }
            else
        #endif
            {
                float3 L = -light.forward;

                // Is it worth sampling the shadow map?
                if ((light.lightDimmer > 0) && (light.shadowDimmer > 0) && // Note: Volumetric can have different dimmer, thus why we test it here
                    dot(bsdfData.normalWS, L) > 0.0)
                {
                    context.shadowValue = GetDirectionalShadowAttenuation(context.shadowContext,
                                                                          posInput.positionSS, posInput.positionWS + L * _ShadowBias, bsdfData.normalWS,
                                                                          light.shadowIndex, L);
                }
            }
        }
    }
    
    PreLightData preLightData = GetPreLightData_UTS(V, posInput, bsdfData);
    
    AggregateLighting aggregateLighting;
    ZERO_INITIALIZE(AggregateLighting, aggregateLighting);
    
    // Evaluate the punctual lights.
    if (featureFlags & LIGHTFEATUREFLAGS_PUNCTUAL)
    {
        uint lightCount, lightStart;

    #ifndef LIGHTLOOP_DISABLE_TILE_AND_CLUSTER
        GetCountAndStart(posInput, LIGHTCATEGORY_PUNCTUAL, lightStart, lightCount);
    #else   // LIGHTLOOP_DISABLE_TILE_AND_CLUSTER
        lightCount = _PunctualLightCount;
        lightStart = 0;
    #endif

        bool fastPath = false;
    #if SCALARIZE_LIGHT_LOOP
        uint lightStartLane0;
        fastPath = IsFastPath(lightStart, lightStartLane0);

        if (fastPath)
        {
            lightStart = lightStartLane0;
        }
    #endif

        // Scalarized loop. All lights that are in a tile/cluster touched by any pixel in the wave are loaded (scalar load), only the one relevant to current thread/pixel are processed.
        // For clarity, the following code will follow the convention: variables starting with s_ are meant to be wave uniform (meant for scalar register),
        // v_ are variables that might have different value for each thread in the wave (meant for vector registers).
        // This will perform more loads than it is supposed to, however, the benefits should offset the downside, especially given that light data accessed should be largely coherent.
        // Note that the above is valid only if wave intriniscs are supported.
        uint v_lightListOffset = 0;
        uint v_lightIdx = lightStart;

        [loop] // vulkan shader compiler can not unroll.
    #if NEED_TO_CHECK_HELPER_LANE
        // On some platform helper lanes don't behave as we'd expect, therefore we prevent them from entering the loop altogether.
        // IMPORTANT! This has implications if ddx/ddy is used on results derived from lighting, however given Lightloop is called in compute we should be
        // sure it will not happen.
        bool isHelperLane = WaveIsHelperLane();
        while (!isHelperLane && v_lightListOffset < lightCount)
    #else
        while (v_lightListOffset < lightCount)
    #endif
        {
            v_lightIdx = FetchIndex(lightStart, v_lightListOffset);
        #if SCALARIZE_LIGHT_LOOP
            uint s_lightIdx = ScalarizeElementIndex(v_lightIdx, fastPath);
        #else
            uint s_lightIdx = v_lightIdx;
        #endif
            if (s_lightIdx == -1)
            {
                break;
            }

            LightData s_lightData = FetchLight(s_lightIdx);

            // If current scalar and vector light index match, we process the light. The v_lightListOffset for current thread is increased.
            // Note that the following should really be ==, however, since helper lanes are not considered by WaveActiveMin, such helper lanes could
            // end up with a unique v_lightIdx value that is smaller than s_lightIdx hence being stuck in a loop. All the active lanes will not have this problem.
            if (s_lightIdx >= v_lightIdx)
            {
                v_lightListOffset++;
                if (IsMatchingLightLayer(s_lightData.lightLayers, builtinData.renderingLayers))
                {
                    DirectLighting lighting = UtsEvaluateBSDF_Punctual(context, posInput, builtinData, s_lightData, bsdfData, preLightData, V, fragInputs.texCoord0.xy);
                    AccumulateDirectLighting(lighting, aggregateLighting);
                }
            }
        }
    }
    
    // Evaluate the directional lights.
    if (featureFlags & LIGHTFEATUREFLAGS_DIRECTIONAL)
    {
        uint i = 0; // Declare once to avoid the D3D11 compiler warning.
        for (i = 0; i < _DirectionalLightCount; ++i)
        {
            if (IsMatchingLightLayer(_DirectionalLightDatas[i].lightLayers, builtinData.renderingLayers))
            {
                DirectLighting lighting = UtsEvaluateBSDF_Directional(context, posInput, builtinData, _DirectionalLightDatas[i], bsdfData, preLightData, V, fragInputs.texCoord0.xy);
                AccumulateDirectLighting(lighting, aggregateLighting);
            }
        }
    }

    // Evaluate the environment lights.
    if (featureFlags & (LIGHTFEATUREFLAGS_ENV | LIGHTFEATUREFLAGS_SKY | LIGHTFEATUREFLAGS_SSREFRACTION | LIGHTFEATUREFLAGS_SSREFLECTION))
    {
        float reflectionHierarchyWeight = 0.0; // Max: 1.0

        uint envLightStart, envLightCount;

        // Fetch first env light to provide the scene proxy for screen space computation
    #ifndef LIGHTLOOP_DISABLE_TILE_AND_CLUSTER
        GetCountAndStart(posInput, LIGHTCATEGORY_ENV, envLightStart, envLightCount);
    #else
        envLightCount = _EnvLightCount;
        envLightStart = 0;
    #endif

        bool fastPath = false;
    #if SCALARIZE_LIGHT_LOOP
        uint envStartFirstLane;
        fastPath = IsFastPath(envLightStart, envStartFirstLane);
    #endif

        // Reflection hierarchy is
        //  1. Screen Space Reflection
        //  2. Environment Reflection
        //  3. Sky Reflection

        // Apply SSR.
    #if (defined(_SURFACE_TYPE_TRANSPARENT) && !defined(_DISABLE_SSR_TRANSPARENT)) || (!defined(_SURFACE_TYPE_TRANSPARENT) && !defined(_DISABLE_SSR))
        {
            IndirectLighting lighting = UtsEvaluateBSDF_ScreenSpaceReflection(posInput, preLightData, reflectionHierarchyWeight);
            AccumulateIndirectLighting(lighting, aggregateLighting);
        }
    #endif

        float3 lightInReflDir = 0.0;
    #ifdef _INDIRECT_DIFFUSE_OFF

    #elif _INDIRECT_DIFFUSE_IBL
        bool replaceBakeDiffuseLighting = false;
        #if !defined(_SURFACE_TYPE_TRANSPARENT) // No SSGI/RTGI/Mixed effect on transparent
        if (_IndirectDiffuseMode != INDIRECTDIFFUSEMODE_OFF)
        {
            replaceBakeDiffuseLighting = true;
        }
    #endif
        
    #if defined(PROBE_VOLUMES_L1) || defined(PROBE_VOLUMES_L2)
        if (!builtinData.isLightmap)
        {
            replaceBakeDiffuseLighting = true;
        }
    #endif

    #if defined(LIGHT_EVALUATION_SKIP_INDIRECT_DIFFUSE)
        replaceBakeDiffuseLighting = false;
    #endif
        
        if (replaceBakeDiffuseLighting)
        {
            UtsEvaluateBSDF_BakeDiffuse(posInput, preLightData, bsdfData, V, builtinData, lightInReflDir);
        }
    #elif _INDIRECT_DIFFUSE_MATCAP
        //builtinData.bakeDiffuseLighting = UtsEvaluateColor_MatCap(posInput.positionWS, bsdfData.normalWS, 0.0);
        UtsEvaluateBSDF_MatCapDiffuse(posInput.positionWS, bsdfData.normalWS, builtinData);
    #elif _INDIRECT_DIFFUSE_RAMP
        UtsEvaluateBSDF_Ramp(posInput, bsdfData, builtinData);
    #endif

        if (featureFlags & LIGHTFEATUREFLAGS_ENV)
        {
        #if _INDIRECT_SPECULAR_OFF
            
        #elif _INDIRECT_SPECULAR_IBL
            context.sampleReflection = SINGLE_PASS_CONTEXT_SAMPLE_REFLECTION_PROBES;

        #if SCALARIZE_LIGHT_LOOP
            if (fastPath)
            {
                envLightStart = envStartFirstLane;
            }
        #endif

            // Scalarized loop, same rationale of the punctual light version
            uint v_envLightListOffset = 0;
            uint v_envLightIdx = envLightStart;
        #if NEED_TO_CHECK_HELPER_LANE
            // On some platform helper lanes don't behave as we'd expect, therefore we prevent them from entering the loop altogether.
            // IMPORTANT! This has implications if ddx/ddy is used on results derived from lighting, however given Lightloop is called in compute we should be
            // sure it will not happen.
            bool isHelperLane = WaveIsHelperLane();
            while (!isHelperLane && v_envLightListOffset < envLightCount)
        #else
            while (v_envLightListOffset < envLightCount)
        #endif
            {
                v_envLightIdx = FetchIndex(envLightStart, v_envLightListOffset);
            #if SCALARIZE_LIGHT_LOOP
                uint s_envLightIdx = ScalarizeElementIndex(v_envLightIdx, fastPath);
            #else
                uint s_envLightIdx = v_envLightIdx;
            #endif
                if (s_envLightIdx == -1)
                {
                    break;
                }

                EnvLightData s_envLightData = FetchEnvLight(s_envLightIdx);

                // If current scalar and vector light index match, we process the light. The v_envLightListOffset for current thread is increased.
                // Note that the following should really be ==, however, since helper lanes are not considered by WaveActiveMin, such helper lanes could
                // end up with a unique v_envLightIdx value that is smaller than s_envLightIdx hence being stuck in a loop. All the active lanes will not have this problem.
                if (s_envLightIdx >= v_envLightIdx)
                {
                    v_envLightListOffset++;
                    if (reflectionHierarchyWeight < 1.0)
                    {
                        if (IsMatchingLightLayer(s_envLightData.lightLayers, builtinData.renderingLayers))
                        {
                            IndirectLighting lighting = UtsEvaluateBSDF_Env(context, posInput, preLightData, s_envLightData, bsdfData, s_envLightData.influenceShapeType, GPUIMAGEBASEDLIGHTINGTYPE_REFLECTION, reflectionHierarchyWeight);
                        #if defined(PROBE_VOLUMES_L1) || defined(PROBE_VOLUMES_L2)

                            if (s_envLightData.normalizeWithAPV > 0 && all(lightInReflDir >= 0))
                            {
                                float factor = GetReflectionProbeNormalizationFactor(lightInReflDir, bsdfData.normalWS, s_envLightData.L0L1, s_envLightData.L2_1, s_envLightData.L2_2);
                                lighting.specularReflected *= factor;
                            }
                        #endif

                            AccumulateIndirectLighting(lighting, aggregateLighting);
                        }
                    }
                    
                }
            }
        #elif _INDIRECT_SPECULAR_MATCAP
            IndirectLighting lighting = UtsEvaluateBSDF_MatCapSpecular(posInput.positionWS, bsdfData, preLightData);
            AccumulateIndirectLighting(lighting, aggregateLighting);
        #endif
        }

    #if _INDIRECT_SPECULAR_IBL
        // Only apply the sky IBL if the sky texture is available
        if ((featureFlags & LIGHTFEATUREFLAGS_SKY) && _EnvLightSkyEnabled)
        {
            // The sky is a single cubemap texture separate from the reflection probe texture array (different resolution and compression)
            context.sampleReflection = SINGLE_PASS_CONTEXT_SAMPLE_SKY;

            // The sky data are generated on the fly so the compiler can optimize the code
            EnvLightData envLightSky = InitSkyEnvLightData(0);

            // Only apply the sky if we haven't yet accumulated enough IBL lighting.
            if (reflectionHierarchyWeight < 1.0)
            {
                IndirectLighting lighting = UtsEvaluateBSDF_Env(context, posInput, preLightData, envLightSky, bsdfData, envLightSky.influenceShapeType, GPUIMAGEBASEDLIGHTINGTYPE_REFLECTION, reflectionHierarchyWeight);
                AccumulateIndirectLighting(lighting, aggregateLighting);
            }
        }
    #endif
    }
    
    if (HasFlag(bsdfData.surfaceFeatures, SURFACEFEATUREFLAGS_ANGEL_RING))
    {
        DirectLighting lighting = UtsEvaluateAngelRing(fragInputs, bsdfData.normalWS, V);
        AccumulateDirectLighting(lighting, aggregateLighting);
    }

    UtsPostEvaluateBSDF(posInput, preLightData, bsdfData, builtinData, aggregateLighting, lightLoopOutput);
}

// UTSLightData GetUTSMainPunctualLightData(BuiltinData builtinData, PositionInputs posInput)
// {
//     UTSLightData mainPunctualLight;

//     uint lightCount, lightStart;
    
// #ifndef LIGHTLOOP_DISABLE_TILE_AND_CLUSTER
//     GetCountAndStart(posInput, LIGHTCATEGORY_PUNCTUAL, lightStart, lightCount);
// #else   // LIGHTLOOP_DISABLE_TILE_AND_CLUSTER
//     lightCount = _PunctualLightCount;
//     lightStart = 0;
// #endif
//     bool fastPath = false;
// #if SCALARIZE_LIGHT_LOOP
//     uint lightStartLane0;
//     fastPath = IsFastPath(lightStart, lightStartLane0);

//     if (fastPath)
//     {
//         lightStart = lightStartLane0;
//     }
// #endif

//     uint v_lightListOffset = 0;
//     uint v_lightIdx = lightStart;
//     float channelAlpha = 0.0f;
//     [loop] // vulkan shader compiler can not unroll.
//     while (v_lightListOffset < lightCount)
//     {
//         v_lightIdx = FetchIndex(lightStart, v_lightListOffset);
// #if SCALARIZE_LIGHT_LOOP
//         uint s_lightIdx = ScalarizeElementIndex(v_lightIdx, fastPath);
// #else
//         uint s_lightIdx = v_lightIdx;
// #endif
//         if (s_lightIdx == -1)
//             break;

//         LightData s_lightData = FetchLight(s_lightIdx);

//         // If current scalar and vector light index match, we process the light. The v_lightListOffset for current thread is increased.
//         // Note that the following should really be ==, however, since helper lanes are not considered by WaveActiveMin, such helper lanes could
//         // end up with a unique v_lightIdx value that is smaller than s_lightIdx hence being stuck in a loop. All the active lanes will not have this problem.
//         if (s_lightIdx >= v_lightIdx)
//         {
//             v_lightListOffset++;
//             if (IsMatchingLightLayer(s_lightData.lightLayers, builtinData.renderingLayers))
//             {
//                 float3 lightDirection;
//                 float4 distances; // {d, d^2, 1/d, d_proj}
//                 GetPunctualLightVectors(posInput.positionWS, s_lightData, lightDirection, distances);
//                 float4 lightColor = EvaluateLight_Punctual(context, posInput, s_lightData, lightDirection, distances);
//                 float3 additionalLightColor = ApplyCurrentExposureMultiplier(lightColor.rgb) * lightColor.a;
//                 const float notDirectional = 1.0f;

//                 UTSLightData utsLightData;
//                 utsLightData.lightColor = additionalLightColor;
//                 utsLightData.lightDirection = lightDirection;
//                 utsLightData.diffuseDimmer = s_lightData.diffuseDimmer;
//                 utsLightData.specularDimmer = s_lightData.specularDimmer;
//                 utsLightData.shadowTint = s_lightData.shadowTint;
//                 utsLightData.penumbraTint = s_lightData.penumbraTint;

//                 if(length(additionalLightColor) >= length(mainPunctualLight.lightColor))
//                 {
//                     mainPunctualLight = utsLightData;
//                 }
//             }
//         }
//     }

//     return mainPunctualLight;
// }

// Todo: calculate the acutal main lighboth dorectional and punctual)t based on the light attenuation, rather than using the main directional light
UTSLightData GetCustomMainLightData(BuiltinData builtinData, UTSLightData mainPunctualLight)
{
    UTSLightData utsLightData;
    int mainLightIndex;

    mainLightIndex = GetUtsMainLightIndex(builtinData);

    if (mainLightIndex == -1 || length(_DirectionalLightDatas[mainLightIndex].color) < length(mainPunctualLight.lightColor))
    {
        utsLightData = mainPunctualLight;
    }
    else
    {
        utsLightData.lightColor = ApplyCurrentExposureMultiplier(_DirectionalLightDatas[mainLightIndex].color);
        utsLightData.lightDirection = -_DirectionalLightDatas[mainLightIndex].forward;
        utsLightData.diffuseDimmer = _DirectionalLightDatas[mainLightIndex].diffuseDimmer;
        utsLightData.specularDimmer = _DirectionalLightDatas[mainLightIndex].specularDimmer;
        utsLightData.shadowTint = _DirectionalLightDatas[mainLightIndex].shadowTint;
        utsLightData.penumbraTint = _DirectionalLightDatas[mainLightIndex].penumbraTint;
    }

    return utsLightData;
}