#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.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"

#pragma kernel KGlobal
#pragma kernel KLocal
#pragma kernel KDynamic
#pragma kernel KPrePass
#pragma kernel KReduction

RWTexture2D<float4> Result;

TEXTURE2D_X(_Source);
TEXTURE2D(_Input);
TEXTURE2D(_Previous);
RW_TEXTURE2D(float, _Current);
RW_TEXTURE2D(float2, _Output);
//TEXTURE2D(_ExposureWeightMask);
//TEXTURE2D(_PreviousExposureTexture);
//TEXTURE2D(_ExposureCurveTexture);

CBUFFER_START(cb)
float4 _Params;
float4 _Params2;
CBUFFER_END

#define ParamMultiplier     _Params.x
#define ParamAverage        _Params.y
#define ParamLdisp          _Params.z
#define ParamDarkAdapt      _Params2.x
#define ParamLightAdapt     _Params2.y
#define ParamRadiusX        _Params2.z
#define ParamRadiusY        _Params2.w
#define ParamSoftness       _Params.w

float logTVIscotopic(float logLuminance)
{
    if (logLuminance <= -3.94)
        return -2.86;
    else if (logLuminance >= -1.44)
        return logLuminance - 0.395;
    else
        return pow(0.405 * logLuminance + 1.6, 2.18) - 2.86;
}

float logTVIphotopic(float logLuminance)
{
    if (logLuminance <= -2.6)
        return -0.72;
    else if (logLuminance >= 1.9)
        return logLuminance - 1.255;
    else
        return pow(0.249 * logLuminance + 0.65, 2.7) - 0.72;
}

float TVImesopic(float logLuminance) {
    float minLum = log10(0.03);
    float maxLum = log10(3);
    float k = pow(clamp((logLuminance - minLum) / (maxLum - minLum), 0.0, 1.0), 1.5);
    float ts = pow(10, logTVIscotopic(logLuminance));
    float tp = pow(10, logTVIphotopic(logLuminance));
    return lerp(ts, tp, k);
}

float2 params(float Lwa, float Lavg)
{
    float logLda = log10(ParamLdisp / 2);
    float logLwa = log10(Lwa);
    float tLda = pow(10, logTVIphotopic(logLda));
    float tLwa = TVImesopic(logLwa);
    float b = 0;
    float m = tLda / tLwa;
    if (Lwa < Lavg)
    {
        float logLavg = log10(Lavg);
        float tLavg = TVImesopic(logLavg);
        m = (tLda / tLavg) * (tLwa / tLavg);
        b = (tLda / tLavg - m) * Lavg;
    }
    return float2(m * ParamMultiplier, b) / ParamLdisp;
}

float adapt(float Lwa, float Lavg)
{
    float t = unity_DeltaTime.x;
    float param = Lavg <= Lwa ? ParamDarkAdapt : ParamLightAdapt;
    return (Lwa - Lavg) * pow(param, t) + Lavg;
}

#define PREPASS_TEX_SIZE 1024.0
#define PREPASS_TEX_HALF_SIZE 512.0

float WeightSample(uint2 pixel, float2 sourceSize)
{
    //float screenDiagonal = 0.25f * (sourceSize.x + sourceSize.y);
    //const float2 kCenter = sourceSize * 0.5f;
    //return 1.0 - saturate(pow(length(kCenter - pixel) / screenDiagonal, ParamSoftness));
    float radius = max(ParamRadiusX, ParamRadiusY);
    float2 ellipseScale = float2(radius / ParamRadiusX, radius / ParamRadiusY) / PREPASS_TEX_HALF_SIZE;
    const float2 kCenter = sourceSize * 0.5f;
    float dist = length(kCenter * ellipseScale - pixel * ellipseScale);
    return saturate(1.0 - pow(abs(dist / radius), ParamSoftness));
}

// https://github.com/needle-mirror/com.unity.render-pipelines.high-definition/blob/5b5be02b68d38e5fc3a38075c94179cf03f87191/Runtime/PostProcessing/Shaders/Exposure.compute#L41-L59
[numthreads(8,8,1)]
void KPrePass (uint2 id : SV_DispatchThreadID)
{
    // For XR, interleave single-pass views in a checkerboard pattern
    UNITY_XR_ASSIGN_VIEW_INDEX((id.x + id.y) % _XRViewCount)

    PositionInputs posInputs = GetPositionInput(float2(id), rcp(PREPASS_TEX_SIZE), uint2(8u, 8u));
    float2 uv = ClampAndScaleUVForBilinear(posInputs.positionNDC);
    float3 color = SAMPLE_TEXTURE2D_X_LOD(_Source, s_linear_clamp_sampler, uv, 0.0).xyz;
    float luma = Luminance(color);

    float weight = WeightSample(id, PREPASS_TEX_SIZE.xx);
    // float logLuma = ComputeEV100FromAvgLuminance(max(luma, 1e-4), MeterCalibrationConstant);

    _Output[posInputs.positionSS] = float2(luma, weight);
}

#define REDUCTION_GROUP_SIZE 16
#define REDUCTION_TOTAL_THREADS 256

groupshared float4 gs_luminances[REDUCTION_TOTAL_THREADS];
groupshared float gs_weights[REDUCTION_TOTAL_THREADS];

// perform reduction, return true iff threadIdx == 0
bool reduce(uint2 groupId, uint2 groupThreadId)
{
    uint threadIdx = groupThreadId.y * REDUCTION_GROUP_SIZE + groupThreadId.x;
    uint2 sampleIdx = (groupId.xy * REDUCTION_GROUP_SIZE + groupThreadId.xy) * 2u;

    // Store 4 pixels & their weights in the lds
    float2 p1 = _Input[sampleIdx + uint2(0u, 0u)].xy;
    float2 p2 = _Input[sampleIdx + uint2(1u, 0u)].xy;
    float2 p3 = _Input[sampleIdx + uint2(0u, 1u)].xy;
    float2 p4 = _Input[sampleIdx + uint2(1u, 1u)].xy;

    float4 smp = float4(p1.x, p2.x, p3.x, p4.x);
    float4 weights = float4(p1.y, p2.y, p3.y, p4.y);

    gs_luminances[threadIdx] = smp * weights;
    gs_weights[threadIdx] = dot(weights, 1.0);

    GroupMemoryBarrierWithGroupSync();

    // Parallel reduction of luminances & weights
    UNITY_UNROLL
        for (uint s = REDUCTION_TOTAL_THREADS / 2u; s > 0u; s >>= 1u)
        {
            if (threadIdx < s)
            {
                gs_luminances[threadIdx] += gs_luminances[threadIdx + s];
                gs_weights[threadIdx] += gs_weights[threadIdx + s];
            }

            GroupMemoryBarrierWithGroupSync();
        }
    
    return threadIdx == 0;
}

// calculate luminance in last step
float2 finish()
{
    float avgLuminance = dot(gs_luminances[0], 0.25);

    if (IsNaN(avgLuminance) || IsInf(avgLuminance))
        avgLuminance = 1.0;

    if (gs_weights[0] > 0.0)
        avgLuminance /= (gs_weights[0] * 0.25);

    return float2(avgLuminance, gs_weights[0]);
}

[numthreads(REDUCTION_GROUP_SIZE, REDUCTION_GROUP_SIZE, 1)]
void KReduction(uint2 groupId : SV_GroupID, uint2 groupThreadId : SV_GroupThreadID)
{
    if (reduce(groupId, groupThreadId)) {
        _Output[groupId.xy] = finish();
    }
}

[numthreads(REDUCTION_GROUP_SIZE, REDUCTION_GROUP_SIZE, 1)]
void KLocal(uint2 groupId : SV_GroupID, uint2 groupThreadId : SV_GroupThreadID)
{
    if (reduce(groupId, groupThreadId)) {
        float avgLuminance = finish().x * ParamMultiplier;
        _Current[groupId.xy] = avgLuminance;
        _Output[groupId.xy] = params(avgLuminance, avgLuminance);
    }
}

[numthreads(REDUCTION_GROUP_SIZE, REDUCTION_GROUP_SIZE, 1)]
void KDynamic(uint2 groupId : SV_GroupID, uint2 groupThreadId : SV_GroupThreadID)
{
    if (reduce(groupId, groupThreadId)) {
        float avgLuminance = finish().x * ParamMultiplier;
        float Lwa = adapt(_Previous[groupId.xy].x, avgLuminance);
        _Current[groupId.xy] = Lwa;
        _Output[groupId.xy] = params(Lwa, avgLuminance);
    }
}

[numthreads(1, 1, 1)]
void KGlobal(uint2 groupId : SV_GroupID, uint2 groupThreadId : SV_GroupThreadID)
{
    float avgLuminance = ParamAverage * ParamMultiplier;
    _Current[groupId.xy] = avgLuminance;
    _Output[groupId.xy] = params(avgLuminance, avgLuminance);
}