//--------------------------------------------------------------------------------------------------
// Definitions
//--------------------------------------------------------------------------------------------------

#pragma kernel VolumetricLightingAllLights VolumetricLighting=VolumetricLightingAllLights LIGHTLOOP_SINGLE_PASS
#pragma kernel VolumetricLightingClustered VolumetricLighting=VolumetricLightingClustered LIGHTLOOP_TILE_PASS   USE_CLUSTERED_LIGHTLIST

#pragma enable_d3d11_debug_symbols

#include "../../../ShaderPass/ShaderPass.cs.hlsl"
#define SHADERPASS    SHADERPASS_VOLUMETRIC_LIGHTING
#define GROUP_SIZE_1D 16
#define GROUP_SIZE_2D (GROUP_SIZE_1D * GROUP_SIZE_1D)
// #define USE_HENYEY_GREENSTEIN_PHASE_FUNCTION

//--------------------------------------------------------------------------------------------------
// Included headers
//--------------------------------------------------------------------------------------------------

#include "../../../../Core/ShaderLibrary/Common.hlsl"
#include "../../../../Core/ShaderLibrary/SpaceFillingCurves.hlsl"
#include "../../../../Core/ShaderLibrary/VolumeRendering.hlsl"

#include "../VolumetricLighting.cs.hlsl"
#define UNITY_MATERIAL_LIT // Need to be defined before including Material.hlsl
#include "../../../ShaderVariables.hlsl"
#include "../../../Lighting/Lighting.hlsl" // This includes Material.hlsl

//--------------------------------------------------------------------------------------------------
// Inputs & outputs
//--------------------------------------------------------------------------------------------------

TEXTURE2D(_DepthTexture);                  // Z-buffer
RW_TEXTURE2D(float4, _CameraColorTexture); // Updated texture

//--------------------------------------------------------------------------------------------------
// Implementation
//--------------------------------------------------------------------------------------------------

struct Ray
{
    float3 originWS;
    float3 directionWS; // Normalized
    float  maxLength;   // In meters
};

// Computes the in-scattered radiance along the ray.
float3 PerformIntegration(PositionInputs posInput, Ray ray, uint numSteps)
{
    float3 scattering = _GlobalFog_Scattering;
    float  extinction = _GlobalFog_Extinction;
#ifdef USE_HENYEY_GREENSTEIN_PHASE_FUNCTION
    float  asymmetry  = _GlobalFog_Asymmetry;
#endif

    LightLoopContext context;
    // ZERO_INITIALIZE(LightLoopContext, context);
    context.shadowContext = InitShadowContext();
    uint featureFlags = 0xFFFFFFFF; // TODO

    float maxDepthVS = posInput.depthVS;

    // Note: we are already using 'unPositionSS' for randomization of LODDitheringTransition().
    float zeta = GenerateHashedRandomFloat(posInput.unPositionSS.yx);

    float du = rcp(numSteps);
    float u0 = 0.25 * du + 0.5 * du * zeta;
    float dt = du * ray.maxLength;
    float t0 = 0;

    float3 radiance = 0;

    for (uint s = 0; s < numSteps; s++)
    {
        float u  = u0 + s * du;          // [0, 1]
        float t  = u * ray.maxLength;    // [0, ray.maxLength]
        float t1 = t + 0.5 * dt;

        float T = TransmittanceIntegralOverHomogeneousInterval(extinction, t0, t1);

        [branch] if (T < 0.0001) break;

        float3 positionWS = ray.originWS + t * ray.directionWS;

        float3 sampleRadiance = 0;

        if (featureFlags & LIGHTFEATUREFLAGS_DIRECTIONAL)
        {
            for (uint i = 0; i < _DirectionalLightCount; ++i)
            {
                // Fetch the light.
                DirectionalLightData lightData = _DirectionalLightDatas[i];

                float3 L         = -lightData.forward; // Lights point backwards in Unity
                float  intensity = 1;
                float3 color     = lightData.color;

            #ifdef USE_HENYEY_GREENSTEIN_PHASE_FUNCTION
                // Note: we apply the scattering coefficient and the constant part of the phase function later.
                intensity *= HenyeyGreensteinPhasePartVarying(asymmetry, dot(L, ray.directionWS));
            #endif

                [branch] if (lightData.shadowIndex >= 0)
                {
                    float shadow = GetDirectionalShadowAttenuation(context.shadowContext, positionWS,
                                   0, lightData.shadowIndex, L, posInput.unPositionSS);

                    intensity *= shadow;
                }

                [branch] if (lightData.cookieIndex >= 0)
                {
                    float3 lightToSample = positionWS - lightData.positionWS;
                    float4 cookie = EvaluateCookie_Directional(context, lightData, lightToSample);

                    color     *= cookie.rgb;
                    intensity *= cookie.a;
                }

                // Compute the amount of in-scattered radiance.
                sampleRadiance += color * intensity;
            }
        }

        if (featureFlags & LIGHTFEATUREFLAGS_PUNCTUAL)
        {
            uint punctualLightCount;

        #ifdef LIGHTLOOP_TILE_PASS
            uint punctualLightStart;

            posInput.depthVS = u * maxDepthVS;
            GetCountAndStart(posInput, LIGHTCATEGORY_PUNCTUAL, punctualLightStart, punctualLightCount);
        #else
            punctualLightCount = _PunctualLightCount;
        #endif

            for (uint i = 0; i < punctualLightCount; ++i)
            {
            #ifdef LIGHTLOOP_TILE_PASS
                uint punctualLightIndex = FetchIndex(punctualLightStart, i);
            #else
                uint punctualLightIndex = i;
            #endif

                // Fetch the light.
                LightData lightData  = _LightDatas[punctualLightIndex];
                int       lightType  = lightData.lightType;

                float3 lightToSample = positionWS - lightData.positionWS;
                float  distSq        = dot(lightToSample, lightToSample);
                float  dist          = sqrt(distSq);
                float3 L             = lightToSample * -rsqrt(distSq);
                float  intensity     = GetPunctualShapeAttenuation(lightData, L, distSq);
                float3 color         = lightData.color;

            #ifdef USE_HENYEY_GREENSTEIN_PHASE_FUNCTION
                // Note: we apply the scattering coefficient and the constant part of the phase function later.
                intensity *= HenyeyGreensteinPhasePartVarying(asymmetry, dot(L, ray.directionWS));
            #endif
                intensity *= Transmittance(OpticalDepthHomogeneous(extinction, dist));

                [branch] if (lightData.shadowIndex >= 0)
                {
                    // TODO: make projector lights cast shadows.
                    float3 offset = 0; // GetShadowPosOffset(nDotL, normal);

                    float shadow = GetPunctualShadowAttenuation(context.shadowContext, positionWS + offset,
                                   0, lightData.shadowIndex, float4(L, dist), posInput.unPositionSS);

                    intensity *= lerp(1, shadow, lightData.shadowDimmer);
                }

                // Projector lights always have a cookies, so we can perform clipping inside the if().
                [branch] if (lightData.cookieIndex >= 0)
                {
                    float4 cookie = EvaluateCookie_Punctual(context, lightData, lightToSample);

                    color     *= cookie.rgb;
                    intensity *= cookie.a;
                }

                // Compute the amount of in-scattered radiance.
                sampleRadiance += color * intensity;
            }
        }

        radiance += sampleRadiance * T;
        t0 = t1;
    }
#ifdef USE_HENYEY_GREENSTEIN_PHASE_FUNCTION
    float3 phaseConstant = scattering * HenyeyGreensteinPhasePartConstant(asymmetry);
#else
    float3 phaseConstant = scattering * IsotropicPhaseFunction();
#endif
    return radiance * phaseConstant;
}

[numthreads(GROUP_SIZE_2D, 1, 1)]
void VolumetricLighting(uint2 groupId       : SV_GroupID,
                        uint  groupThreadId : SV_GroupThreadID)
{
    // Note: any factor of 64 is a suitable wave size for our algorithm.
    uint waveIndex = groupThreadId / 64;
    uint laneIndex = groupThreadId % 64;
    uint quadIndex = laneIndex / 4;

    // Arrange threads in the Morton order to optimally match the memory layout of GCN tiles.
    uint  mortonCode = groupThreadId;
    uint2 localCoord = DecodeMorton2D(mortonCode);
    uint2 tileAnchor = groupId * GROUP_SIZE_1D;
    uint2 pixelCoord = tileAnchor + localCoord;
    uint2 tileCoord  = pixelCoord / GetTileSize();

    if (pixelCoord.x >= (uint)_ScreenSize.x || pixelCoord.y >= (uint)_ScreenSize.y) { return; }

    // Idea: zenith angle based distance limiting to simulate aerial perspective?
#ifdef UNITY_REVERSED_Z
    float z = max(LOAD_TEXTURE2D(_DepthTexture, pixelCoord).r, 0 + 0.001);
#else
    float z = min(LOAD_TEXTURE2D(_DepthTexture, pixelCoord).r, 1 - 0.001);
#endif

    PositionInputs posInput = GetPositionInput(pixelCoord, _ScreenSize.zw, tileCoord);
    UpdatePositionInput(z, _InvViewProjMatrix, _ViewProjMatrix, posInput);

    Ray cameraRay;

    // Note: the camera ray does not start on the the near (camera sensor) plane.
    // While this is not correct (strictly speaking), the introduced error is small.
    cameraRay.originWS     = GetCurrentViewPosition();
    cameraRay.directionWS  = posInput.positionWS - cameraRay.originWS;
    cameraRay.maxLength    =  sqrt(dot(cameraRay.directionWS, cameraRay.directionWS));
    cameraRay.directionWS *= rsqrt(dot(cameraRay.directionWS, cameraRay.directionWS)); // Normalize

    float rayT = Transmittance(OpticalDepthHomogeneous(_GlobalFog_Extinction, cameraRay.maxLength));

    const int numSamples = 64;
    float3 inL = PerformIntegration(posInput, cameraRay, numSamples);

    // In-place UAV updates do not work on Intel GPUs.
    _CameraColorTexture[pixelCoord] = float4(rayT * _CameraColorTexture[pixelCoord].rgb + inL, _CameraColorTexture[pixelCoord].a);
}