#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch

#pragma multi_compile TONEMAPPING_NONE TONEMAPPING_NEUTRAL TONEMAPPING_ACES_APPROX TONEMAPPING_ACES_FULL TONEMAPPING_CUSTOM TONEMAPPING_EXTERNAL
#pragma multi_compile_local _ HDR_COLORSPACE_CONVERSION
#pragma multi_compile_local GRADE_IN_SRGB GRADE_IN_ACESCG

#if (defined(TONEMAPPING_ACES_APPROX) || defined(TONEMAPPING_ACES_FULL))
#define TONEMAPPING_ACES
#endif
#define TONEMAPPING_USE_FULL_ACES defined(TONEMAPPING_ACES_FULL)

#pragma kernel KBuild

#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/ACES.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Color.hlsl"
#if defined(HDR_COLORSPACE_CONVERSION)
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/HDROutput.hlsl"
#endif

TEXTURE3D(_LogLut3D);
RW_TEXTURE3D(float4, _OutputTexture);

TEXTURE2D(_CurveMaster);
TEXTURE2D(_CurveRed);
TEXTURE2D(_CurveGreen);
TEXTURE2D(_CurveBlue);
TEXTURE2D(_CurveHueVsHue);
TEXTURE2D(_CurveHueVsSat);
TEXTURE2D(_CurveSatVsSat);
TEXTURE2D(_CurveLumVsSat);

SAMPLER(sampler_LogLut3D);
SAMPLER(sampler_LinearClamp);

CBUFFER_START(cb0)
    float4 _Size;               // x: lut_size, y: 1 / (lut_size - 1), zw: unused
    float4 _LogLut3D_Params;    // x: 1 / lut_size, y: lut_size - 1, z: contribution, w: unused
    float4 _ColorBalance;       // xyz: LMS coeffs, w: unused
    float4 _ColorFilter;        // xyz: color, w: unused
    float4 _ChannelMixerRed;    // xyz: rgb coeffs, w: unused
    float4 _ChannelMixerGreen;  // xyz: rgb coeffs, w: unused
    float4 _ChannelMixerBlue;   // xyz: rgb coeffs, w: unused
    float4 _HueSatCon;          // x: hue shift, y: saturation, z: contrast, w: unused
    float4 _Lift;               // xyz: color, w: unused
    float4 _Gamma;              // xyz: color, w: unused
    float4 _Gain;               // xyz: color, w: unused
    float4 _Shadows;            // xyz: color, w: unused
    float4 _Midtones;           // xyz: color, w: unused
    float4 _Highlights;         // xyz: color, w: unused
    float4 _ShaHiLimits;        // xy: shadows min/max, zw: highlight min/max
    float4 _SplitShadows;       // xyz: color, w: balance
    float4 _SplitHighlights;    // xyz: color, w: unused
    float4 _Params;             // x: enable grading, yzw: unused

    // Custom tonemapping settings
    float4 _CustomToneCurve;
    float4 _ToeSegmentA;
    float4 _ToeSegmentB;
    float4 _MidSegmentA;
    float4 _MidSegmentB;
    float4 _ShoSegmentA;
    float4 _ShoSegmentB;

    float4 _HDROutputParams;
    float4 _HDROutputParams2;
    #define _MinNits    _HDROutputParams.x
    #define _MaxNits    _HDROutputParams.y
    #define _PaperWhite _HDROutputParams.z
    #define _RangeReductionMode    (int)_HDROutputParams2.x
    #define _HueShift    _HDROutputParams2.y

CBUFFER_END

float GetLuminance(float3 colorLinear)
{
    #if defined(TONEMAPPING_ACES) || defined(GRADE_IN_ACESCG)
    return AcesLuminance(colorLinear);
    #else
    return Luminance(colorLinear);
    #endif
}

float EvaluateCurve(TEXTURE2D(curve), float t)
{
    float x = SAMPLE_TEXTURE2D_LOD(curve, sampler_LinearClamp, float2(t, 0.0), 0.0).x;
    return saturate(x);
}

float3 RotateToColorGradingSpace_LogToLinear(float3 logColor)
{
#if defined(TONEMAPPING_ACES) || defined(GRADE_IN_ACESCG)
    return ACES_to_ACEScg(ACEScc_to_ACES(half3(logColor)));
#elif defined(GRADE_IN_SRGB)
    return LogCToLinear(logColor);
#endif
}

float3 RotateToColorGradingSpace_LinearToLog(float3 linearColor)
{
#if defined(TONEMAPPING_ACES) || defined(GRADE_IN_ACESCG)
    return ACES_to_ACEScc(unity_to_ACES(half3(linearColor)));
#elif defined(GRADE_IN_SRGB)
    return LinearToLogC(linearColor);
#endif
}


float3 RotateToColorGradeOutputSpace(float3 gradedColor)
{
#ifdef TONEMAPPING_ACES

    // In ACES workflow we return graded color in ACEScg, we move to ACES (AP0) later on
    return gradedColor;

#elif defined(HDR_COLORSPACE_CONVERSION) // HDR but not ACES workflow

    // If we are doing HDR we expect grading to finish at Rec2020. Any supplemental rotation is done inside the various options.
#ifdef GRADE_IN_ACESCG
    return ACEScg_to_Rec2020(half3(gradedColor));
#elif defined(GRADE_IN_SRGB)
    return RotateRec709ToRec2020(gradedColor);
#endif

#else // Nor ACES or HDR

#ifdef GRADE_IN_ACESCG
    // If we are not HDR or ACES, we move back to Rec709
    return ACEScg_to_unity(half3(gradedColor));
#elif defined(GRADE_IN_SRGB)
    // We already graded in sRGB
    return gradedColor;
#endif

#endif
}

// Note: when the ACES tonemapper is selected the grading steps will be done using ACES spaces
float3 ColorGrade(float3 colorLutSpace)
{
    // Switch back to linear
#ifdef HDR_COLORSPACE_CONVERSION
    // For LUT purposes, we assume 1.0 == 100 nits (max PQ value is 10k nits)
    float3 colorLinear = PQToLinear(colorLutSpace, 100.0f);
#else
    float3 colorLinear = LogCToLinear(colorLutSpace);
#endif

    // White balance in LMS space
    float3 colorLMS = LinearToLMS(colorLinear);
    colorLMS *= _ColorBalance.xyz;
    colorLinear = LMSToLinear(colorLMS);

    // Do contrast in log after white balance
    float3 colorLog = RotateToColorGradingSpace_LinearToLog(colorLinear);

    colorLog = (colorLog - ACEScc_MIDGRAY) * _HueSatCon.z + ACEScc_MIDGRAY;

    colorLinear = RotateToColorGradingSpace_LogToLinear(colorLog);

    // Color filter is just an unclipped multiplier
    colorLinear *= _ColorFilter.xyz;

    // Do NOT feed negative values to the following color ops
    colorLinear = max(0.0, colorLinear);

    // Split toning
    // As counter-intuitive as it is, to make split-toning work the same way it does in Adobe
    // products we have to do all the maths in gamma-space...
    float balance = _SplitShadows.w;
    float3 colorGamma = PositivePow(colorLinear, 1.0 / 2.2);

    float luma = saturate(GetLuminance(saturate(colorGamma)) + balance);
    float3 splitShadows = lerp((0.5).xxx, _SplitShadows.xyz, 1.0 - luma);
    float3 splitHighlights = lerp((0.5).xxx, _SplitHighlights.xyz, luma);
    colorGamma = SoftLight(colorGamma, splitShadows);
    colorGamma = SoftLight(colorGamma, splitHighlights);

    colorLinear = PositivePow(colorGamma, 2.2);

    // Channel mixing (Adobe style)
    colorLinear = float3(
        dot(colorLinear, _ChannelMixerRed.xyz),
        dot(colorLinear, _ChannelMixerGreen.xyz),
        dot(colorLinear, _ChannelMixerBlue.xyz)
    );

    // Shadows, midtones, highlights
    luma = GetLuminance(colorLinear);
    float shadowsFactor = 1.0 - smoothstep(_ShaHiLimits.x, _ShaHiLimits.y, luma);
    float highlightsFactor = smoothstep(_ShaHiLimits.z, _ShaHiLimits.w, luma);
    float midtonesFactor = 1.0 - shadowsFactor - highlightsFactor;
    colorLinear = colorLinear * _Shadows.xyz * shadowsFactor
                + colorLinear * _Midtones.xyz * midtonesFactor
                + colorLinear * _Highlights.xyz * highlightsFactor;

    // Lift, gamma, gain
    colorLinear = colorLinear * _Gain.xyz + _Lift.xyz;
    colorLinear = sign(colorLinear) * pow(abs(colorLinear), _Gamma.xyz);

    // HSV operations
    float satMult;
    float3 hsv = RgbToHsv(colorLinear);
    {
        // Hue Vs Sat
        satMult = EvaluateCurve(_CurveHueVsSat, hsv.x) * 2.0;

        // Sat Vs Sat
        satMult *= EvaluateCurve(_CurveSatVsSat, hsv.y) * 2.0;

        // Lum Vs Sat
        satMult *= EvaluateCurve(_CurveLumVsSat, Luminance(colorLinear)) * 2.0;

        // Hue Shift & Hue Vs Hue
        float hue = hsv.x + _HueSatCon.x;
        float offset = EvaluateCurve(_CurveHueVsHue, hue) - 0.5;
        hue += offset;
        hsv.x = RotateHue(hue, 0.0, 1.0);
    }
    colorLinear = HsvToRgb(hsv);

    // Global saturation
    luma = GetLuminance(colorLinear);
    colorLinear = luma.xxx + (_HueSatCon.yyy * satMult) * (colorLinear - luma.xxx);

    // YRGB curves
    // Conceptually these need to be in range [0;1] and from an artist-workflow perspective it's
    // easier to deal with
    colorLinear = FastTonemap(colorLinear);
    {
        const float kHalfPixel = (1.0 / 128.0) / 2.0;
        float3 c = colorLinear;

        // Y (master)
        c += kHalfPixel.xxx;
        float mr = EvaluateCurve(_CurveMaster, c.r);
        float mg = EvaluateCurve(_CurveMaster, c.g);
        float mb = EvaluateCurve(_CurveMaster, c.b);
        c = float3(mr, mg, mb);

        // RGB
        c += kHalfPixel.xxx;
        float r = EvaluateCurve(_CurveRed, c.r);
        float g = EvaluateCurve(_CurveGreen, c.g);
        float b = EvaluateCurve(_CurveBlue, c.b);
        colorLinear = float3(r, g, b);
    }
    colorLinear = FastTonemapInvert(colorLinear);

    colorLinear = max(0.0, colorLinear);

    return RotateToColorGradeOutputSpace(colorLinear);
}

// Used for debugging - see the ColorGrading option in FrameSettings
float3 NeutralColorGrade(float3 colorLutSpace)
{
    // Switch back to linear
#ifdef HDR_COLORSPACE_CONVERSION
    // For LUT purposes, we assume 1.0 == 100 nits (max PQ value is 10k nits)
    float3 colorLinear = PQToLinear(colorLutSpace, 100.0f);
#ifdef GRADE_IN_ACESCG
    // In ACESCG grading mode color grading space is ACEScg so RotateToColorGradeOutputSpace will move out of this space. Here that means we need to move into that space to balance it.
    colorLinear = unity_to_ACEScg(half3(colorLinear));
#endif
#else
    float3 colorLinear = LogCToLinear(colorLutSpace);
#endif

    return RotateToColorGradeOutputSpace(colorLinear);
}

float3 Tonemap(float3 colorLinear)
{
    #if defined(TONEMAPPING_NEUTRAL)
    {
        colorLinear = NeutralTonemap(colorLinear);
    }
    #elif defined(TONEMAPPING_CUSTOM)
    {
        colorLinear = CustomTonemap(colorLinear, _CustomToneCurve.xyz, _ToeSegmentA, _ToeSegmentB.xy, _MidSegmentA, _MidSegmentB.xy, _ShoSegmentA, _ShoSegmentB.xy);
    }
    #elif defined(TONEMAPPING_ACES)
    {
        // Note: input is actually ACEScg (AP1 w/ linear encoding)
        float3 aces = ACEScg_to_ACES(colorLinear);
        colorLinear = AcesTonemap(aces);
    }
    #elif defined(TONEMAPPING_EXTERNAL)
    {
        float3 colorLutSpace = saturate(LinearToLogC(colorLinear));
        float3 colorLut = ApplyLut3D(TEXTURE3D_ARGS(_LogLut3D, sampler_LogLut3D), colorLutSpace, _LogLut3D_Params.xy);
        colorLinear = lerp(colorLinear, colorLut, _LogLut3D_Params.z);
    }
    #endif

    return colorLinear;
}


float3 ProcessColorForHDR(float3 gradedColor)
{
#ifdef HDR_COLORSPACE_CONVERSION

#ifdef TONEMAPPING_ACES
    gradedColor = ACEScg_to_ACES(gradedColor);
    return HDRMappingACES(gradedColor.rgb, _PaperWhite, _MinNits, _MaxNits, _RangeReductionMode, true);
#else
    return HDRMappingFromRec2020(gradedColor.rgb, _PaperWhite, _MinNits, _MaxNits, _RangeReductionMode, _HueShift, true);
#endif

#endif
    return gradedColor;
}


// Note: according to the specs the maximum thread group size for Metal/Desktop is 1024. 8x8x8 is
// 512 so it shouldn't be an issue... except with some Intel chipsets where for some reason it won't
// allow anything higher than 256 threads. We'll use 4x4x4 then.
// Ultimately it would nice to expose `maxTotalThreadsPerThreadgroup` for Metal...
// Source: https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf
// It is important to keep this in sync with the group-size declared in PostProcessSystem.cs
[numthreads(4,4,4)]
void KBuild(uint3 dispatchThreadId : SV_DispatchThreadID)
{
    // Lut space
    // We use Alexa LogC (El 1000) to store the LUT as it provides a good enough range (~58.85666)
    // and is good enough to be stored in fp16 without losing precision in the darks
    // If HDR Output is enabled, then the LUT space is PQ.
    float3 colorLutSpace = float3(dispatchThreadId) * _Size.yyy;

    // Color grade & tonemap
    float3 gradedColor = colorLutSpace;

    if (_Params.x > 0.0)
    {
        gradedColor = ColorGrade(gradedColor);
    }
    else
    {
        // Skip grading
        gradedColor = NeutralColorGrade(gradedColor);
    }

    // Negative colors are technically allowed under ACEScg but we often use 11-11-10 anyway,
    // so clamping is fine and it helps avoiding edge cases for tonemapping
    gradedColor = max(gradedColor, 0.0);

#ifdef HDR_COLORSPACE_CONVERSION
    gradedColor = ProcessColorForHDR(gradedColor);
#else
    gradedColor = Tonemap(gradedColor);
#endif

    _OutputTexture[dispatchThreadId] = float4(max(gradedColor, 0.0), 1.0);
}