#include <metal_stdlib>
using namespace metal;

struct GrainParams {
    uint width;
    uint height;
    uint frameSeed;
    float amount;    // overall grain intensity, 0..1
    float size;      // grain clump size multiplier, >0 (1 = baseline)
    float halation;  // warm glow around highlights, 0..1
    float halationThreshold; // 0..1, luminance level above which halation begins
    float bloom;     // soft overall highlight lift, 0..1
    float grainSaturation; // 0..1, 0 = monochrome grain, 1 = full-color grain
    float gateAspect; // film gate width:height aspect ratio, 0 = no gate
    float gateGlowSpread; // 0..1, how far the halation/bloom glow spreads into the gate mask
    float gateOpacity; // 0..1, opacity of the masked-out gate area (1 = solid black)
    float gateCornerRadius; // 0..1, rounds the gate aperture's corners
    float gateCornerTangent; // 0..1, biases the corner's tangent point between its two edges (0.5 = circular)
    float gateBlur; // 0..1, widens the gate edge transition to soften/feather the gate boundary
    uint grainAffectsGate; // 0 = grain only on image, 1 = grain also visible over the gate mask
};

// Integer-lattice hash (no sin()) - avoids the periodic aliasing that
// produces visible diagonal weave patterns when sampled at non-integer
// or low frequencies.
inline float hash21(uint x, uint y, uint seed, uint channel) {
    uint h = x * 374761393u + y * 668265263u + seed * 2654435761u + channel * 2246822519u;
    h = (h ^ (h >> 13)) * 1274126177u;
    h ^= (h >> 16);
    return float(h & 0x00FFFFFFu) / float(0x01000000u);
}

// Multi-scale grain: a coarse "clump" layer (silver-halide crystal
// clusters) plus a fine per-pixel layer (individual crystal sparkle),
// each generated independently per color channel so the three layers
// don't move in lockstep like a printed texture.
inline float generateGrain(float2 pos, float invSize, uint frameSeed, uint channel) {
    // Coarse clump layer: sampled on an irregularly-jittered coarse grid.
    float2 coarsePos = pos * invSize;
    uint2 cell = uint2(floor(coarsePos));
    float2 frac2 = coarsePos - float2(cell);

    float n00 = hash21(cell.x,     cell.y,     frameSeed, channel);
    float n10 = hash21(cell.x + 1, cell.y,     frameSeed, channel);
    float n01 = hash21(cell.x,     cell.y + 1, frameSeed, channel);
    float n11 = hash21(cell.x + 1, cell.y + 1, frameSeed, channel);

    // Smoothstep interpolation gives soft, organic clump edges rather
    // than hard cell boundaries.
    float2 sf = frac2 * frac2 * (3.0 - 2.0 * frac2);
    float clump = mix(mix(n00, n10, sf.x), mix(n01, n11, sf.x), sf.y);

    // Fine per-pixel layer: full-resolution independent noise, always at
    // native pixel pitch regardless of "size" so grain never looks like a
    // smooth blur even at large clump sizes.
    uint2 px = uint2(pos);
    float fine = hash21(px.x, px.y, frameSeed * 9781u + 17u, channel);

    // Blend: larger clumps (invSize small) lean more on the coarse layer,
    // finer grain leans more on the per-pixel layer.
    float fineWeight = clamp(invSize, 0.15, 0.85);
    return mix(clump, fine, fineWeight);
}

inline float softScale(float value) {
    if (value >= 0.0 && value <= 1.0) {
        return value;
    }
    if (value < 0.0) {
        return value * 0.5 - value * 0.5 * value;
    }
    float excess = value - 1.0;
    return 1.0 - (excess * 0.5) / (1.0 + excess * 0.5);
}

kernel void grainKernel(
    device const float* src [[ buffer(0) ]],
    device float* dst [[ buffer(1) ]],
    constant GrainParams& params [[ buffer(2) ]],
    uint2 gid [[ thread_position_in_grid ]])
{
    if (gid.x >= params.width || gid.y >= params.height) {
        return;
    }

    uint index = (gid.y * params.width + gid.x) * 3;

    float r = src[index + 0];
    float g = src[index + 1];
    float b = src[index + 2];

    float luminance = 0.2126 * r + 0.7152 * g + 0.0722 * b;
    float luminanceForWeight = clamp(luminance, 0.0, 1.0);

    float weight = 4.0 * luminanceForWeight * (1.0 - luminanceForWeight);

    // Larger "size" lowers the clump sampling frequency, producing
    // bigger, coarser grain clumps.
    float invSize = 1.0 / max(params.size, 0.01) * 0.5;
    float2 pos = float2(gid.x, gid.y);

    // Independent noise per color channel - real film stock has separate
    // dye/silver layers, so R/G/B grain isn't perfectly correlated.
    float noiseR = generateGrain(pos, invSize, params.frameSeed, 0u);
    float noiseG = generateGrain(pos, invSize, params.frameSeed, 1u);
    float noiseB = generateGrain(pos, invSize, params.frameSeed, 2u);

    // Grain Saturation: blend each channel's noise toward the luma of the
    // three (their average), so at 0 all three channels share the same
    // noise value (monochrome grain) and at 1 they vary independently
    // (full-color grain).
    float noiseLuma = (noiseR + noiseG + noiseB) / 3.0;
    noiseR = mix(noiseLuma, noiseR, params.grainSaturation);
    noiseG = mix(noiseLuma, noiseG, params.grainSaturation);
    noiseB = mix(noiseLuma, noiseB, params.grainSaturation);

    float noiseOffsetR = (noiseR - 0.5) * weight * params.amount;
    float noiseOffsetG = (noiseG - 0.5) * weight * params.amount;
    float noiseOffsetB = (noiseB - 0.5) * weight * params.amount;

    // Halation: warm orange-red glow that bleeds out of bright areas.
    // Halation Threshold restricts the glow to highlights above the given
    // luminance level - 0 behaves like the old always-on curve, higher
    // values squeeze the glow into only the brightest areas.
    float halationRange = max(1.0 - params.halationThreshold, 0.001);
    float halationLevel = clamp((luminanceForWeight - params.halationThreshold) / halationRange, 0.0, 1.0);
    float halationGlow = pow(halationLevel, 2.0) * params.halation;

    // Bloom: soft overall lift applied to bright areas across all channels.
    float bloomGlow = pow(luminanceForWeight, 2.0) * params.bloom;

    float outR = r + noiseOffsetR + halationGlow + bloomGlow;
    float outG = g + noiseOffsetG + halationGlow * 0.4 + bloomGlow;
    float outB = b + noiseOffsetB + bloomGlow;

    // Film Gate: letterbox/pillarbox the frame to the chosen camera
    // format's aperture aspect ratio. Pixels outside the gate are masked
    // toward black; the masked band picks up a warm halation/bloom glow
    // bleeding in from the gate edge, as on a real film gate.
    if (params.gateAspect > 0.0001) {
        float frameAspect = float(params.width) / float(params.height);
        float2 uv = (pos + 0.5) / float2(params.width, params.height);
        float2 centered = uv - 0.5;

        float2 halfExtent = (params.gateAspect > frameAspect)
            ? float2(0.5, 0.5 * frameAspect / params.gateAspect)
            : float2(0.5 * params.gateAspect / frameAspect, 0.5);

        // Rounded-rectangle signed distance: gateCornerRadius (0..1) maps
        // to a fraction of the gate's shorter half-extent so the corners
        // can be rounded all the way to a stadium/ellipse shape at 100%.
        // Gate Corner Tangent (0..1, 0.5 = circular) biases where the
        // corner is tangent to its two edges, stretching the corner radius
        // into an ellipse toward the x or y edge as it moves off-center.
        // Gate Corner Tangent (0..1, 0.5 = circular) blends the corner
        // curve's superellipse exponent: low values pull it toward a
        // diamond/chamfer, high values push it toward a squared-off
        // "squircle", while the radius slider keeps controlling overall
        // corner size.
        float radius = params.gateCornerRadius * min(halfExtent.x, halfExtent.y);
        float superellipseN = (params.gateCornerTangent <= 0.5)
            ? mix(1.0, 2.0, params.gateCornerTangent * 2.0)
            : mix(2.0, 8.0, (params.gateCornerTangent - 0.5) * 2.0);
        float2 innerHalfExtent = halfExtent - radius;
        float2 q = abs(centered) - innerHalfExtent;
        float2 qPos = max(q, 0.0);
        float cornerLen = (radius > 1e-6)
            ? pow(pow(qPos.x, superellipseN) + pow(qPos.y, superellipseN), 1.0 / superellipseN)
            : length(qPos);
        float dist = min(max(q.x, q.y), 0.0) + cornerLen - radius;

        float edgeWidth = mix(0.0015, 0.08, params.gateBlur);
        float gateMask = 1.0 - smoothstep(0.0, edgeWidth, dist);

        // Warm glow bleeding from the gate edge into the masked area,
        // driven by the Halation/Bloom sliders. Gate Glow Spread controls
        // how quickly the glow decays with distance from the edge - low
        // spread keeps it tight to the edge, high spread lets it bleed
        // further into the masked-out area.
        float decay = mix(140.0, 12.0, params.gateGlowSpread);
        float glow = exp(-max(dist, 0.0) * decay) * (params.halation * 0.6 + params.bloom * 0.6);

        // Grain Affects Gate: if enabled, the grain noise already added to
        // outR/outG/outB carries through into the masked-out area too;
        // otherwise the masked-out area is built from the clean (pre-grain)
        // pixel so grain only appears over the visible image.
        float3 outsideBase = (params.grainAffectsGate != 0)
            ? float3(outR, outG, outB)
            : float3(r, g, b);
        // Gate Opacity: 1 = solid black gate (real film gate look), 0 =
        // the underlying image shows through with just the edge glow.
        outsideBase = mix(outsideBase, float3(0.0), params.gateOpacity);
        float3 outside = outsideBase + float3(glow, glow * 0.4, 0.0);

        outR = mix(outside.r, outR, gateMask);
        outG = mix(outside.g, outG, gateMask);
        outB = mix(outside.b, outB, gateMask);
    }

    dst[index + 0] = clamp(outR, 0.0, 1.0);
    dst[index + 1] = clamp(outG, 0.0, 1.0);
    dst[index + 2] = clamp(outB, 0.0, 1.0);
}
