// ExposureMatchKernel.metal
//
// Metal equivalent of ExposureMatchKernel.cu:
//  - AI-calculated exposure applied as a base-2 (stops) multiplier
//  - 3x3 "mood" and "skin protect" color matrices applied to the result
//  - Per-pixel skin mask blends between the two graded outputs
//
// Buffer layout: row-major, 32-bit float, RGB interleaved (3 floats/pixel).

#include <metal_stdlib>
using namespace metal;

struct ExposureMatchParams {
    uint width;
    uint height;
    float exposureStops; // AI-calculated exposure, in stops
    float slope;         // sensitometric slope (gamma) trim, -1..1 (0 = no change)
    float contrast;      // pivot contrast trim, -1..1 (0 = no change)
    float saturation;     // saturation trim, -1..1 (0 = no change)
    float highlightRolloff; // 0..1, shoulder compression strength (0 = none)
    float shadowKnee;       // 0..1, toe compression strength (0 = none)
    // Chroma Pulp: per-zone tint offsets (-0.5..0.5, 0 = neutral)
    float pulpShadowR, pulpShadowG, pulpShadowB;
    float pulpMidR, pulpMidG, pulpMidB;
    float pulpHighlightR, pulpHighlightG, pulpHighlightB;
    // Chroma Pulp: luma-range pivots (0..1) and tangent/softness (0..1) per zone
    float pulpShadowRange, pulpShadowTangent;
    float pulpHighlightRange, pulpHighlightTangent;
};

kernel void exposureMatchKernel(
    const device float* inBuffer        [[ buffer(0) ]],
    device float* outBuffer             [[ buffer(1) ]],
    const device float* skinMask        [[ buffer(2) ]], // [height * width], 0..1
    constant ExposureMatchParams& params [[ buffer(3) ]],
    constant float3x3& moodMatrix       [[ buffer(4) ]], // background grade
    constant float3x3& skinMatrix       [[ buffer(5) ]], // skin tone protect grade
    uint2 gid                            [[ thread_position_in_grid ]])
{
    if (gid.x >= params.width || gid.y >= params.height) {
        return;
    }

    const uint pixelIdx = gid.y * params.width + gid.x;
    const uint base = pixelIdx * 3;

    // 1) Exposure: multiply by 2^stops.
    const float exposureFactor = exp2(params.exposureStops);

    float3 rgb = float3(inBuffer[base + 0],
                         inBuffer[base + 1],
                         inBuffer[base + 2]) * exposureFactor;

    // 1b) Slope: sensitometric gamma trim. The photographic characteristic
    //     curve is D = gamma * log10(H) + const, where gamma is the slope
    //     of the straight-line portion (Basic Photographic Sensitometry
    //     Workbook). In linear-light terms this is a power-law about a
    //     pivot: out = pivot * (in/pivot)^gamma. gamma = 2^slope so
    //     slope in [-1, 1] maps to gamma in [0.5, 2.0] (flatter <-> steeper).
    const float kSlopePivot = 0.18f;
    const float gamma = exp2(params.slope);
    rgb = kSlopePivot * pow(max(rgb / kSlopePivot, 0.0f), gamma);

    // 1c) Contrast: pivot around 18% grey, so the trim only steepens or
    //     flattens the curve without shifting overall density.
    const float kContrastPivot = 0.18f;
    rgb = (rgb - kContrastPivot) * (1.0f + params.contrast) + kContrastPivot;

    // 1d) Highlight Rolloff (Shoulder): emulates Kodak print-film highlight
    //     compression - blends toward an exponential shoulder above a knee
    //     point so highlights compress gently into the white point rather
    //     than clipping abruptly. 0 = no change (identity).
    const float kShoulderPoint = 0.65f;
    if (params.highlightRolloff > 0.0f) {
        float3 shoulderOut = kShoulderPoint + (1.0f - kShoulderPoint) *
            (1.0f - exp(-(rgb - kShoulderPoint) * 4.0f));
        rgb = select(rgb, mix(rgb, shoulderOut, params.highlightRolloff), rgb > kShoulderPoint);
    }

    // 1e) Shadow Knee (Toe): emulates Kodak negative toe compression -
    //     blends toward a power-law toe below a knee point so shadow
    //     detail rolls off gently instead of crushing to black.
    //     0 = no change (identity).
    const float kToePoint = 0.1f;
    if (params.shadowKnee > 0.0f) {
        float n = mix(1.0f, 0.3f, params.shadowKnee);
        float3 toeOut = kToePoint * pow(max(rgb, 0.0f) / kToePoint, float3(n));
        rgb = select(rgb, toeOut, rgb < kToePoint);
    }

    // 2) Apply both 3x3 color matrices ("polynomial" first-order term)
    //    to the exposed pixel. Matrices are passed row-major; Metal's
    //    float3x3 is column-major, so we multiply with the vector on
    //    the right using the transposed convention (matrix * vector
    //    where matrix columns correspond to source channels).
    const float3 moodRGB = moodMatrix * rgb;
    const float3 skinRGB = skinMatrix * rgb;

    // 3) Blend mood vs. skin-protect grade using the per-pixel skin mask
    //    weight (1 = fully skin-protected, 0 = fully mood grade).
    const float w = skinMask[pixelIdx];

    float3 blended = mix(moodRGB, skinRGB, w);

    // 4) Saturation: blend toward (or away from) the pixel's luma.
    const float luma = dot(blended, float3(0.2126f, 0.7152f, 0.0722f));
    blended = mix(float3(luma), blended, 1.0f + params.saturation);

    // 4b) Chroma Pulp: 3-way shadows/mids/highlights color tint, selected
    //     by per-pixel luma. The shadow/highlight range pivots define
    //     where each zone is fully active; the tangent (softness) values
    //     widen the smoothstep transition centered on each pivot.
    {
        float pulpLuma = dot(blended, float3(0.2126f, 0.7152f, 0.0722f));

        float shadowSoft = max(params.pulpShadowTangent, 0.001f);
        float highlightSoft = max(params.pulpHighlightTangent, 0.001f);

        float wShadow = 1.0f - smoothstep(params.pulpShadowRange - shadowSoft,
                                           params.pulpShadowRange + shadowSoft, pulpLuma);
        float wHighlight = smoothstep(params.pulpHighlightRange - highlightSoft,
                                       params.pulpHighlightRange + highlightSoft, pulpLuma);
        wShadow = clamp(wShadow, 0.0f, 1.0f);
        wHighlight = clamp(wHighlight, 0.0f, 1.0f - wShadow);
        float wMid = 1.0f - wShadow - wHighlight;

        float3 shadowTint = float3(params.pulpShadowR, params.pulpShadowG, params.pulpShadowB);
        float3 midTint = float3(params.pulpMidR, params.pulpMidG, params.pulpMidB);
        float3 highlightTint = float3(params.pulpHighlightR, params.pulpHighlightG, params.pulpHighlightB);

        blended += shadowTint * wShadow + midTint * wMid + highlightTint * wHighlight;
    }

    const float3 outRGB = clamp(blended, 0.0f, 1.0f);

    outBuffer[base + 0] = outRGB.x;
    outBuffer[base + 1] = outRGB.y;
    outBuffer[base + 2] = outRGB.z;
}
