#include "Packages/com.unity.render-pipelines.high-definition/Runtime/PostProcessing/Shaders/MotionBlurCommon.hlsl"

#pragma kernel MotionVecPreppingCS MOTION_VEC_PREPPING

#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch


#define SKIP_PREPPING_IF_NOT_NEEDED defined(PLATFORM_SUPPORTS_WAVE_INTRINSICS)

// Special clamps for camera component.
#pragma multi_compile  NO_SPECIAL_CLAMP CAMERA_ROT_CLAMP CAMERA_TRANS_CLAMP CAMERA_FULL_CLAMP CAMERA_SEPARATE_CLAMP CAMERA_DISABLE_CAMERA

// We use polar coordinates. This has the advantage of storing the length separately and we'll need the length several times.
float2 EncodeMotionVector(float2 motionVec)
{
    float motionVecLen = length(motionVec);
    if (motionVecLen < 0.0001)
    {
        return 0.0;
    }
    else
    {
        float theta = atan2(motionVec.y, motionVec.x)  * (0.5 / PI) + 0.5;
        return float2(motionVecLen, theta);
    }
}

float2 ClampMotionVec(float2 motionVec, float maxMotionVec)
{
    float len = length(motionVec);
    float newLen = min(len, maxMotionVec);
    return (len > 1e-4f && newLen > 1e-4f) ? newLen * (motionVec * rcp(len)) : 0.0;
}


float2 GetDeltaNDCVec(float4 positionWS, float4 prevPosWS, float4x4 currM, float4x4 prevM)
{
    float4 clipWP = mul(currM, positionWS);
    float4 clipPrevWP = mul(prevM, prevPosWS);

    clipWP.xy /= clipWP.w;
    clipPrevWP.xy /= clipPrevWP.w;

    float2 outDeltaVec = (clipWP.xy - clipPrevWP.xy);
    outDeltaVec *= 0.5f;

#if UNITY_UV_STARTS_AT_TOP
    outDeltaVec.y = -outDeltaVec.y;
#endif

    return outDeltaVec;
}


// Prep motion vectors so that the velocity due to rotation is clamped more lightly
// A bit of code duplication but keeping separate make it clearer.

#ifdef NO_SPECIAL_CLAMP

float2 ComputeMotionVec(PositionInputs posInput, float2 sampledMotionVec)
{
    return ClampMotionVec(sampledMotionVec * _MotionBlurIntensity, _MotionBlurMaxMotionVec);
}

#elif defined(CAMERA_DISABLE_CAMERA)

float2 ComputeMotionVec(PositionInputs posInput, float2 sampledMotionVec)
{
    float4 worldPos = float4(posInput.positionWS, 1.0);
    float4 prevPos = worldPos;
    float2 cameraMv = GetDeltaNDCVec(worldPos, prevPos, UNITY_MATRIX_UNJITTERED_VP, UNITY_MATRIX_PREV_VP);

    return ClampMotionVec((sampledMotionVec - cameraMv) * _MotionBlurIntensity, _MotionBlurMaxMotionVec);
}

#elif defined(CAMERA_TRANS_CLAMP) || defined(CAMERA_SEPARATE_CLAMP)

float2 ComputeMotionVec(PositionInputs posInput, float2 sampledMotionVec)
{
    float4 worldPos = float4(posInput.positionWS, 1.0);
    float4 prevPos = worldPos;

    // Calculate translation part
    float2 clampedCameraTranslationMV = 0;
    float2 cameraTranslationMv = 0;
    {
        // Note: potentially wrong if projection matrix changes, but should be rare enough and will last only one frame.
        cameraTranslationMv = GetDeltaNDCVec(worldPos, worldPos - float4(_PrevCamPosRWS.xyz, 0.0), UNITY_MATRIX_UNJITTERED_VP, UNITY_MATRIX_UNJITTERED_VP);
        clampedCameraTranslationMV = ClampMotionVec(cameraTranslationMv * _MotionBlurIntensity, _CameraTranslationClampNDC);
    }

    float2 clampedCameraRotationMV = 0;
    float2 cameraRotationMv = 0;
#if defined(CAMERA_SEPARATE_CLAMP)
    {
        cameraRotationMv = GetDeltaNDCVec(worldPos, worldPos, _CurrVPMatrixNoTranslation, _PrevVPMatrixNoTranslation);
        clampedCameraRotationMV = ClampMotionVec(cameraRotationMv * _MotionBlurIntensity, _CameraRotationClampNDC);
    }
#endif

    float2 mvWithoutCameraComponents = sampledMotionVec - cameraRotationMv - cameraTranslationMv;
    mvWithoutCameraComponents = ClampMotionVec(mvWithoutCameraComponents * _MotionBlurIntensity, _MotionBlurMaxMotionVec);
    return mvWithoutCameraComponents + clampedCameraTranslationMV + clampedCameraRotationMV;
}


#elif defined(CAMERA_ROT_CLAMP) || defined(CAMERA_FULL_CLAMP)

#if defined(CAMERA_ROT_CLAMP)
#define _CameraClampThreshold _CameraRotationClampNDC
#else
#define _CameraClampThreshold _CameraFullClampNDC
#endif

float2 ComputeMotionVec(PositionInputs posInput, float2 sampledMotionVec)
{
    float4 worldPos = float4(posInput.positionWS, 1.0);
    float4 prevPos = worldPos;

    float4x4 prevVP = UNITY_MATRIX_PREV_VP;
    float4x4 currVP = UNITY_MATRIX_UNJITTERED_VP;

#if defined(CAMERA_ROT_CLAMP)
    prevVP = _PrevVPMatrixNoTranslation;
    currVP = _CurrVPMatrixNoTranslation;
#endif

    float2 cameraMv = GetDeltaNDCVec(worldPos, prevPos, currVP, prevVP);

    float2 velocityWithoutCameraComponent = sampledMotionVec - cameraMv;
    cameraMv *= _MotionBlurIntensity;
    float2 clampedCameraMotionVec = ClampMotionVec(cameraMv, _CameraClampThreshold);
    return ClampMotionVec(velocityWithoutCameraComponent * _MotionBlurIntensity, _MotionBlurMaxMotionVec) + clampedCameraMotionVec;
}

#else

#error

#endif

[numthreads(8, 8,1)]
void MotionVecPreppingCS(uint3 dispatchThreadId : SV_DispatchThreadID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);

    float3 motionVecAndDepth = 0.0f;

    float2 motionVecUV = FromOutputPosSSToPreupsampleUV(dispatchThreadId.xy);
    motionVecUV = ClampAndScaleUVForPoint(motionVecUV);

    float4 motionVecBufferSample = SAMPLE_TEXTURE2D_X_LOD(_CameraMotionVectorsTexture, s_point_clamp_sampler, motionVecUV, 0);
    // if we have a value > 1.0f, it means we have selected the "no motion option", hence we force motionVec 0.
    bool forceNoMotion = PixelSetAsNoMotionVectors(motionVecBufferSample);

    float2 motionVec;
    DecodeMotionVector(motionVecBufferSample, motionVec);

    float depth = LoadCameraDepth(FromOutputPosSSToPreupsamplePosSS(dispatchThreadId.xy));


    if ( !forceNoMotion
#if SKIP_PREPPING_IF_NOT_NEEDED
        && WaveActiveAnyTrue(dot(motionVec, motionVec) * _ScreenMagnitudeSq > _MinMotionVecThresholdSq)
#endif
        )
    {

        PositionInputs posInput = GetPositionInput(dispatchThreadId.xy, _PostProcessScreenSize.zw, depth, UNITY_MATRIX_I_VP, UNITY_MATRIX_V);

        float2 finalMotionVec = ComputeMotionVec(posInput, motionVec);
        motionVecAndDepth.xy = EncodeMotionVector(finalMotionVec);
        motionVecAndDepth.z = posInput.linearDepth;
    }
    else
    {
        motionVecAndDepth.z = LinearEyeDepth(depth, _ZBufferParams);
    }

    _MotionVecAndDepth[COORD_TEXTURE2D_X(dispatchThreadId.xy)] = motionVecAndDepth;
}