static const int kReferenceCellIndex = 0;
static const int kMetadataCellIndex = 1;

float2 cellCenterPixelForIndex(int index)
{
    return float2(1.0 + float(index) * 3.0, 1.0);
}

float2 cellCenterUvForIndex(ShaderContext context, int index)
{
    return (cellCenterPixelForIndex(index) + 0.5) / context.outputResolution;
}

bool pixelIsInsideCell(float2 pixelCoord, int index)
{
    float minX = float(index) * 3.0;
    float maxX = minX + 3.0;
    return pixelCoord.x >= minX && pixelCoord.x < maxX && pixelCoord.y >= 0.0 && pixelCoord.y < 3.0;
}

float4 readStoredCell(ShaderContext context, int index)
{
    if (context.feedbackAvailable <= 0)
        return float4(0.0, 0.0, 0.0, 0.0);
    return sampleFeedback(cellCenterUvForIndex(context, index));
}

float3 sampleProbeAverage(ShaderContext context)
{
    float2 clampedSize = clamp(sampleSize, float2(0.001, 0.001), float2(1.0, 1.0));
    float2 halfSize = clampedSize * 0.5;
    float2 minUv = clamp(sampleCenter - halfSize, float2(0.0, 0.0), float2(1.0, 1.0));
    float2 maxUv = clamp(sampleCenter + halfSize, float2(0.0, 0.0), float2(1.0, 1.0));

    float3 total = float3(0.0, 0.0, 0.0);
    float weight = 0.0;
    for (int y = 0; y < 3; ++y)
    {
        for (int x = 0; x < 3; ++x)
        {
            float2 t = float2((float(x) + 0.5) / 3.0, (float(y) + 0.5) / 3.0);
            float2 uv = lerp(minUv, maxUv, t);
            total += sampleLayerInput(uv).rgb;
            weight += 1.0;
        }
    }

    return total / max(weight, 0.0001);
}

float4 storeReferenceState(ShaderContext context)
{
    float2 pixelCoord = floor(context.uv * context.outputResolution);
    float3 currentSample = sampleProbeAverage(context);

    float previousTriggerCount = context.feedbackAvailable > 0
        ? readStoredCell(context, kMetadataCellIndex).r
        : -1.0;
    float currentTriggerCount = float(captureReference);
    bool captureNow = context.feedbackAvailable <= 0 || currentTriggerCount > previousTriggerCount + 0.5;

    float3 storedReference = context.feedbackAvailable > 0
        ? readStoredCell(context, kReferenceCellIndex).rgb
        : currentSample;
    if (captureNow)
        storedReference = currentSample;

    float4 metadata = float4(currentTriggerCount, captureReferenceTime, 0.0, 1.0);
    if (!captureNow && context.feedbackAvailable > 0)
        metadata = readStoredCell(context, kMetadataCellIndex);

    if (pixelIsInsideCell(pixelCoord, kReferenceCellIndex))
        return float4(storedReference, 1.0);
    if (pixelIsInsideCell(pixelCoord, kMetadataCellIndex))
        return metadata;

    return float4(0.0, 0.0, 0.0, 1.0);
}

float rectMask(float2 uv, float2 minUv, float2 maxUv)
{
    if (uv.x < minUv.x || uv.x > maxUv.x)
        return 0.0;
    if (uv.y < minUv.y || uv.y > maxUv.y)
        return 0.0;
    return 1.0;
}

float borderMask(float2 uv, float2 minUv, float2 maxUv, float thickness)
{
    float outer = rectMask(uv, minUv, maxUv);
    float inner = rectMask(uv, minUv + thickness, maxUv - thickness);
    return saturate(outer - inner);
}

float luminance(float3 color)
{
    return dot(color, float3(0.2126, 0.7152, 0.0722));
}

float3 activeReferenceColor(float3 capturedReference)
{
    // Enum parameters are exposed as their zero-based option index.
    // 0 = captured sample, 1 = manual color.
    return referenceSource == 1 ? manualReference.rgb : capturedReference;
}

float4 displayReferenceCompare(ShaderContext context)
{
    float3 liveColor = sampleLayerInput(context.uv).rgb;
    float3 currentSample = sampleProbeAverage(context);
    float3 capturedReference = sampleVideo(cellCenterUvForIndex(context, kReferenceCellIndex)).rgb;
    float3 storedReference = activeReferenceColor(capturedReference);
    float3 delta = currentSample - storedReference;
    float3 absoluteDelta = abs(delta);
    float differenceMagnitude = max(absoluteDelta.r, max(absoluteDelta.g, absoluteDelta.b));

    float3 displayColor = liveColor;
    float opacity = saturate(overlayOpacity);

    float2 halfSize = clamp(sampleSize, float2(0.001, 0.001), float2(1.0, 1.0)) * 0.5;
    float2 boxMin = clamp(sampleCenter - halfSize, float2(0.0, 0.0), float2(1.0, 1.0));
    float2 boxMax = clamp(sampleCenter + halfSize, float2(0.0, 0.0), float2(1.0, 1.0));
    float pixelThickness = 2.0 / max(min(context.outputResolution.x, context.outputResolution.y), 1.0);
    float outerOutline = borderMask(context.uv, boxMin - pixelThickness, boxMax + pixelThickness, pixelThickness);
    float innerOutline = borderMask(context.uv, boxMin, boxMax, pixelThickness);
    if (outerOutline > 0.5)
        displayColor = float3(0.0, 0.0, 0.0);
    if (innerOutline > 0.5)
        displayColor = lerp(displayColor, float3(1.0, 0.0, 0.0), opacity);

    float2 swatchSize = float2(0.06, 0.07);
    float2 panelOrigin = float2(0.03, 0.04);
    float2 gap = float2(0.075, 0.0);
    float2 refMin = panelOrigin;
    float2 curMin = panelOrigin + gap;
    float2 diffMin = panelOrigin + gap * 2.0;
    float2 refMax = refMin + swatchSize;
    float2 curMax = curMin + swatchSize;
    float2 diffMax = diffMin + swatchSize;
    float swatchBorder = min(swatchSize.x, swatchSize.y) * 0.08;

    float refFill = rectMask(context.uv, refMin, refMax);
    float curFill = rectMask(context.uv, curMin, curMax);
    float diffFill = rectMask(context.uv, diffMin, diffMax);
    float refOutline = borderMask(context.uv, refMin, refMax, swatchBorder);
    float curOutline = borderMask(context.uv, curMin, curMax, swatchBorder);
    float diffOutline = borderMask(context.uv, diffMin, diffMax, swatchBorder);

    if (refFill > 0.5)
        displayColor = storedReference;
    if (curFill > 0.5)
        displayColor = currentSample;
    if (diffFill > 0.5)
    {
        float3 neutralBase = float3(0.5, 0.5, 0.5);
        float3 signedDeltaDisplay = saturate(neutralBase + delta * max(differenceGain, 0.0) * 0.5);
        displayColor = signedDeltaDisplay;
    }
    if (refOutline > 0.5 || curOutline > 0.5 || diffOutline > 0.5)
        displayColor = float3(0.0, 0.0, 0.0);

    // Approximate the difference in two operator-friendly axes:
    // warm/cool leans red versus blue, and green/magenta leans green versus
    // the average of red and blue. Centered bars make "match" obvious.
    float warmCool = clamp((delta.r - delta.b) * max(differenceGain, 0.0), -1.0, 1.0);
    float greenMagenta = clamp((delta.g - (delta.r + delta.b) * 0.5) * max(differenceGain, 0.0), -1.0, 1.0);
    float brightnessDelta = clamp((luminance(currentSample) - luminance(storedReference)) * max(differenceGain, 0.0) * 1.5, -1.0, 1.0);

    float barWidth = 0.18;
    float barHeight = 0.018;
    float halfBarWidth = barWidth * 0.5;
    float2 warmCoolMin = float2(0.03, 0.13);
    float2 warmCoolMax = warmCoolMin + float2(barWidth, barHeight);
    float2 tintMin = float2(0.03, 0.157);
    float2 tintMax = tintMin + float2(barWidth, barHeight);
    float2 brightnessMin = float2(0.03, 0.184);
    float2 brightnessMax = brightnessMin + float2(barWidth, barHeight);
    float centerX = warmCoolMin.x + halfBarWidth;

    float warmCoolFill = rectMask(context.uv, warmCoolMin, warmCoolMax);
    float tintFill = rectMask(context.uv, tintMin, tintMax);
    float brightnessFill = rectMask(context.uv, brightnessMin, brightnessMax);
    float warmCoolOutline = borderMask(context.uv, warmCoolMin, warmCoolMax, barHeight * 0.12);
    float tintOutline = borderMask(context.uv, tintMin, tintMax, barHeight * 0.12);
    float brightnessOutline = borderMask(context.uv, brightnessMin, brightnessMax, barHeight * 0.12);
    float targetHalfWidth = 0.0015;
    float warmCoolCenter = rectMask(context.uv, float2(centerX - targetHalfWidth, warmCoolMin.y), float2(centerX + targetHalfWidth, warmCoolMax.y));
    float tintCenter = rectMask(context.uv, float2(centerX - targetHalfWidth, tintMin.y), float2(centerX + targetHalfWidth, tintMax.y));
    float brightnessCenter = rectMask(context.uv, float2(centerX - targetHalfWidth, brightnessMin.y), float2(centerX + targetHalfWidth, brightnessMax.y));

    float warmCoolPosition = centerX + warmCool * halfBarWidth;
    float tintPosition = centerX + greenMagenta * halfBarWidth;
    float brightnessPosition = centerX + brightnessDelta * halfBarWidth;
    float indicatorHalfWidth = 0.0018;
    float warmCoolIndicator = rectMask(context.uv, float2(warmCoolPosition - indicatorHalfWidth, warmCoolMin.y), float2(warmCoolPosition + indicatorHalfWidth, warmCoolMax.y));
    float tintIndicator = rectMask(context.uv, float2(tintPosition - indicatorHalfWidth, tintMin.y), float2(tintPosition + indicatorHalfWidth, tintMax.y));
    float brightnessIndicator = rectMask(context.uv, float2(brightnessPosition - indicatorHalfWidth, brightnessMin.y), float2(brightnessPosition + indicatorHalfWidth, brightnessMax.y));

    if (warmCoolFill > 0.5)
    {
        float gradientT = saturate((context.uv.x - warmCoolMin.x) / max(barWidth, 0.0001));
        float3 coolColor = float3(0.18, 0.52, 1.0);
        float3 warmColor = float3(1.0, 0.48, 0.10);
        float3 gradientColor = gradientT <= 0.5
            ? lerp(coolColor, float3(1.0, 1.0, 1.0), gradientT * 2.0)
            : lerp(float3(1.0, 1.0, 1.0), warmColor, (gradientT - 0.5) * 2.0);
        displayColor = gradientColor;
    }
    if (tintFill > 0.5)
    {
        float gradientT = saturate((context.uv.x - tintMin.x) / max(barWidth, 0.0001));
        float3 magentaColor = float3(1.0, 0.25, 0.75);
        float3 greenColor = float3(0.18, 0.92, 0.32);
        float3 gradientColor = gradientT <= 0.5
            ? lerp(magentaColor, float3(1.0, 1.0, 1.0), gradientT * 2.0)
            : lerp(float3(1.0, 1.0, 1.0), greenColor, (gradientT - 0.5) * 2.0);
        displayColor = gradientColor;
    }
    if (brightnessFill > 0.5)
    {
        float gradientT = saturate((context.uv.x - brightnessMin.x) / max(barWidth, 0.0001));
        float3 darkColor = float3(0.18, 0.18, 0.18);
        float3 gradientColor = lerp(darkColor, float3(1.0, 1.0, 1.0), gradientT);
        displayColor = gradientColor;
    }
    if (warmCoolCenter > 0.5 || tintCenter > 0.5 || brightnessCenter > 0.5)
        displayColor = float3(0.12, 0.45, 1.0);
    if (warmCoolIndicator > 0.5 || tintIndicator > 0.5 || brightnessIndicator > 0.5)
        displayColor = float3(1.0, 0.0, 0.0);
    if (warmCoolOutline > 0.5 || tintOutline > 0.5 || brightnessOutline > 0.5)
        displayColor = float3(0.0, 0.0, 0.0);

    return float4(saturate(displayColor), 1.0);
}