// We need N bounces given that we want to support complex light paths
#pragma max_recursion_depth 11

// HDRP include
#define SHADER_TARGET 50

#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Color.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Sampling/Sampling.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/Builtin/BuiltinData.hlsl"

// Ray tracing includes
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/ShaderVariablesRaytracing.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/Common/AtmosphericScatteringRayTracing.hlsl"

// We need this for the potential volumetric integration on camera misses
#define HAS_LIGHTLOOP

// Path tracing includes
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/PathTracing/Shaders/PathTracingIntersection.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/PathTracing/Shaders/PathTracingSampling.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/PathTracing/Shaders/PathTracingVolume.hlsl"

// Input(s)
float4x4 _PixelCoordToViewDirWS;
int      _RaytracingCameraSkyEnabled;
float4   _RaytracingCameraClearColor;

// DoF related parameters
float4   _PathTracedDoFConstants; // x: aperture radius, y: focus distance, zw: unused

// Output(s)
RWTexture2D<float4> _RadianceTexture;

[shader("miss")]
void MissCamera(inout PathIntersection pathIntersection : SV_RayPayload)
{
    bool skyEnabled = _EnvLightSkyEnabled && _RaytracingCameraSkyEnabled;
    pathIntersection.value = skyEnabled ? SampleSkyTexture(WorldRayDirection(), 0.0, 0).xyz : _RaytracingCameraClearColor.xyz * GetInverseCurrentExposureMultiplier();
    pathIntersection.alpha = skyEnabled ? 1.0 : _RaytracingCameraClearColor.w;

    ApplyFogAttenuation(WorldRayOrigin(), WorldRayDirection(), pathIntersection.value);

    if (_EnableVolumetricFog && _RaytracingMinRecursion <= 1)
    {
        float3 envValue = pathIntersection.value;

        // Generate a 4D unit-square sample for this depth, from our QMC sequence
        float4 inputSample = GetSample4D(pathIntersection.pixelCoord, _RaytracingSampleIndex, 0);

        // Compute volumetric scattering
        pathIntersection.value = 0.0;
        pathIntersection.t = FLT_MAX;
        float pdf = 1.0;
        bool sampleLocalLights;
        if (SampleVolumeScatteringPosition(inputSample.w, pathIntersection.t, pdf, sampleLocalLights))
        {
            ComputeVolumeScattering(pathIntersection, inputSample.xyz, sampleLocalLights);

            // Apply the pdf
            pathIntersection.value /= pdf;

            // Apply volumetric attenuation
            ApplyFogAttenuation(WorldRayOrigin(), WorldRayDirection(), pathIntersection.t, pathIntersection.value);
        }

        // Reinject the environment value
        pathIntersection.value += envValue;
    }
}

[shader("miss")]
void MissLight(inout PathIntersection pathIntersection : SV_RayPayload)
{
}

[shader("miss")]
void MissMaterial(inout PathIntersection pathIntersection : SV_RayPayload)
{
    if ((_RaytracingMaxRecursion - pathIntersection.remainingDepth) < _RaytracingMinRecursion)
    {
        pathIntersection.value = 0.0;
        return;
    }

    pathIntersection.value = _EnvLightSkyEnabled ? SampleSkyTexture(WorldRayDirection(), 0.0, 0).xyz : 0.0;
    ApplyFogAttenuation(WorldRayOrigin(), WorldRayDirection(), pathIntersection.value);
}

[shader("raygeneration")]
void RayGen()
{
    uint2 LaunchIndex = DispatchRaysIndex().xy;

    // Get the current pixel coordinates
    uint2 currentPixelCoord = uint2(LaunchIndex.x, LaunchIndex.y);

    // Jitter them (we use 4x10 dimensions of our sequence during path tracing atm, so pick the next available ones)
    float3 jitteredPixelCoord = float3(currentPixelCoord, 1.0);
    jitteredPixelCoord.x += GetSample(currentPixelCoord, _RaytracingSampleIndex, 40);
    jitteredPixelCoord.y += GetSample(currentPixelCoord, _RaytracingSampleIndex, 41);

    // Compute the ray direction from those coordinates (for zero aperture)
    float3 directionWS = -normalize(mul(jitteredPixelCoord, (float3x3)_PixelCoordToViewDirWS));
    float3 cameraPosWS = GetPrimaryCameraPosition();

    float apertureRadius = _PathTracedDoFConstants.x;
    if (apertureRadius > 0.0)
    {
        // Compute the ray origin and direction for a lens with non-zero aperture

        // Sample the lens apperture using the next available dimensions (we use 40 for path tracing, 2 for sub-pixel jittering, 64 for SSS -> 106, 107)
        float r1 = GetSample(currentPixelCoord, _RaytracingSampleIndex, 106);
        float r2 = GetSample(currentPixelCoord, _RaytracingSampleIndex, 107);
        float2 uv = apertureRadius * SampleDiskUniform(r1, r2);

        // Compute the new ray origin ( _ViewMatrix[0] = right, _ViewMatrix[1] = up, _ViewMatrix[2] = forward )
        cameraPosWS += _ViewMatrix[0].xyz * uv.x + _ViewMatrix[1].xyz * uv.y;

        // Compute the focus point by intersecting the pinhole ray with the focus plane
        float focusDistance = _PathTracedDoFConstants.y;
        float t = focusDistance / dot(directionWS, _ViewMatrix[2].xyz);
        float3 focusPointWS = GetPrimaryCameraPosition() - t * directionWS;

        // The new ray direction should pass through the focus point
        directionWS = normalize(focusPointWS - cameraPosWS);
    }

    // Create the ray descriptor for this pixel
    RayDesc rayDescriptor;
    rayDescriptor.Origin = cameraPosWS;
    rayDescriptor.Direction = directionWS;
    rayDescriptor.TMin = _RaytracingCameraNearPlane;
    rayDescriptor.TMax = FLT_INF;

    // Create and init the PathIntersection structure for this
    PathIntersection pathIntersection;
    pathIntersection.value = 1.0;
    pathIntersection.alpha = 1.0;
    pathIntersection.remainingDepth = _RaytracingMaxRecursion;
    pathIntersection.pixelCoord = currentPixelCoord;
    pathIntersection.maxRoughness = 0.0;

    // In order to achieve filtering for the textures, we need to compute the spread angle of the pixel
    pathIntersection.cone.spreadAngle = _RaytracingPixelSpreadAngle;
    pathIntersection.cone.width = 0.0;

    // Evaluate the ray intersection
    TraceRay(_RaytracingAccelerationStructure, RAY_FLAG_CULL_BACK_FACING_TRIANGLES, RAYTRACINGRENDERERFLAG_PATH_TRACING, 0, 1, 0, rayDescriptor, pathIntersection);

    _RadianceTexture[currentPixelCoord] = float4(pathIntersection.value, pathIntersection.alpha);
}

// This should never be called, return magenta just in case
[shader("closesthit")]
void ClosestHit(inout PathIntersection pathIntersection : SV_RayPayload, AttributeData attributeData : SV_IntersectionAttributes)
{
    pathIntersection.value = float3(1.0, 0.0, 0.5);
}