//===----------------------------------------------------------------------===//
//  Spatial.metal — Edge-adaptive spatial upsampling, in Metal Shading
//  Language for CSR.
//
//  Two self-contained passes:
//
//    1. EASU  (Edge Adaptive Spatial Upsampling)  — the actual upscale. Reads
//       a 12-tap neighbourhood, estimates the local edge direction and how
//       "edge-like" the region is, then resamples with a Lanczos-like kernel
//       that is rotated onto the edge and squashed across it. That is what
//       keeps diagonals clean instead of staircased.
//
//    2. RCAS  (Robust Contrast Adaptive Sharpening) — a follow-up sharpen that
//       only ever *subtracts* from a 4-neighbour ring, with per-pixel limits
//       that guarantee it can never push a pixel outside the local range. So
//       it sharpens without ringing or clipping.
//
//  Written from the published algorithm. The original
//  uses hand-rolled bit-hack approximations for rcp/rsqrt because it targets
//  GPUs where those are slow. Apple GPUs have fast native reciprocal and
//  inverse-sqrt units, so this port calls them directly — same math, less code,
//  and slightly *more* accurate than the reference.
//
//  COLOR SPACE: these kernels operate on perceptual (gamma-encoded,
//  non-linear) color, not linear light. Both kernels are therefore fed a
//  texture view with a plain (non-`_sRGB`) pixel format, so the sampler hands
//  back the stored gamma-encoded values untouched. See MetalTextureFactory.
//===----------------------------------------------------------------------===//

#include <metal_stdlib>
using namespace metal;

// Parameters shared with the Swift side. Must stay layout-identical to the
// `SpatialUniforms` struct in SpatialUpscaler.swift.
struct SpatialUniforms {
    float2 inputSize;      // source texture dimensions, pixels
    float2 outputSize;     // destination texture dimensions, pixels
    float2 invInputSize;   // 1.0 / inputSize, precomputed on the CPU
    float2 scale;          // inputSize / outputSize
    float  sharpness;      // RCAS sharpness, 0 (max sharp) .. 2 (soft)
    float  _pad0;
    float2 _pad1;
};

//===----------------------------------------------------------------------===//
//  Shared helpers
//===----------------------------------------------------------------------===//

// The edge detector.s luma proxy. Deliberately *not* a correct luminance: green-heavy
// and cheap, which is all the edge detector needs.
static inline float edgeLuma(float3 c) {
    return c.g + 0.5f * (c.r + c.b);
}

static inline float max3f(float a, float b, float c) { return max(a, max(b, c)); }
static inline float min3f(float a, float b, float c) { return min(a, min(b, c)); }

// Reciprocal that yields 0 instead of inf for a 0 input. The reference relies
// on its approximation saturating; we have to be explicit.
static inline float safeRcp(float x) {
    return x == 0.0f ? 0.0f : 1.0f / x;
}

//===----------------------------------------------------------------------===//
//  Pass 1 — EASU
//===----------------------------------------------------------------------===//

// Accumulates the edge direction and edge "strength" from one 5-tap cross,
// weighted by how close the output pixel sits to that cross's centre.
//
// The cross is laid out as:
//        a
//      b c d
//        e
//
// `w` is the bilinear weight of this corner. EASU runs this four times, once
// per corner of the 2x2 the output pixel falls inside, so the direction
// estimate varies smoothly across the source pixel instead of snapping.
static inline void easuSet(thread float2 &dir,
                           thread float  &len,
                           float2 pp,
                           int    corner,     // 0=TL 1=TR 2=BL 3=BR
                           float  lA, float lB, float lC, float lD, float lE) {
    float w = 0.0f;
    switch (corner) {
        case 0: w = (1.0f - pp.x) * (1.0f - pp.y); break;
        case 1: w =         pp.x  * (1.0f - pp.y); break;
        case 2: w = (1.0f - pp.x) *         pp.y;  break;
        default:w =         pp.x  *         pp.y;  break;
    }

    // Horizontal: gradient across the cross, normalised by the larger of the
    // two one-sided differences. A smooth ramp gives dirX large but lenX small;
    // a hard step gives both large. That separation is the whole trick.
    float dc   = lD - lC;
    float cb   = lC - lB;
    float lenX = safeRcp(max(abs(dc), abs(cb)));
    float dirX = lD - lB;
    dir.x += dirX * w;
    lenX   = saturate(abs(dirX) * lenX);
    len   += lenX * lenX * w;

    // Vertical: same, on the other axis.
    float ec   = lE - lC;
    float ca   = lC - lA;
    float lenY = safeRcp(max(abs(ec), abs(ca)));
    float dirY = lE - lA;
    dir.y += dirY * w;
    lenY   = saturate(abs(dirY) * lenY);
    len   += lenY * lenY * w;
}

// One resample tap. `off` is the tap's offset from the output pixel measured in
// source pixels; it gets rotated into edge-space and scaled anisotropically
// before the kernel is evaluated, which is what tilts the filter onto the edge.
static inline void easuTap(thread float3 &accColor,
                           thread float  &accWeight,
                           float2 off,
                           float2 dir,
                           float2 len2,
                           float  lob,
                           float  clp,
                           float3 c) {
    // Rotate the offset so +x runs along the detected edge and +y across it.
    float2 v;
    v.x = (off.x *  dir.x) + (off.y * dir.y);
    v.y = (off.x * -dir.y) + (off.y * dir.x);

    // Squash across the edge / stretch along it.
    v *= len2;

    float d2 = min(v.x * v.x + v.y * v.y, clp);

    // Lanczos-2 lookalike, built from two squared quadratics so it needs no
    // sin() and no sqrt():
    //     (25/16 * (2/5 * x^2 - 1)^2 - (25/16 - 1)) * (lob * x^2 - 1)^2
    // The right factor is the adjustable window; `clp` is set so it reaches
    // exactly zero at the window edge.
    float wB = (2.0f / 5.0f) * d2 - 1.0f;
    float wA = lob * d2 - 1.0f;
    wB *= wB;
    wA *= wA;
    wB = (25.0f / 16.0f) * wB - (25.0f / 16.0f - 1.0f);
    float w = wB * wA;

    accColor  += c * w;
    accWeight += w;
}

kernel void csr_easu(texture2d<float, access::sample> src [[texture(0)]],
                      texture2d<float, access::write>  dst [[texture(1)]],
                      constant SpatialUniforms &u             [[buffer(0)]],
                      uint2 gid                            [[thread_position_in_grid]]) {
    if (gid.x >= u.outputSize.x || gid.y >= u.outputSize.y) { return; }

    constexpr sampler pointClamp(coord::pixel, filter::nearest, address::clamp_to_edge);

    // Map this output pixel back to a source-pixel coordinate. The -0.5 shift
    // moves us from pixel-corner to pixel-centre convention.
    float2 pp = (float2(gid) + 0.5f) * u.scale - 0.5f;
    float2 fp = floor(pp);
    pp -= fp;   // fractional position inside the source pixel, [0,1)

    // The 12-tap neighbourhood: a 4x4 block with its corners dropped.
    //
    //        b  c
    //     e  f  g  h
    //     i  j  k  l
    //        n  o
    //
    // `f` is the source pixel we landed in; `pp` is measured from f.
    #define TAP(dx, dy) src.sample(pointClamp, fp + float2(dx, dy) + 0.5f).rgb

    float3 bC = TAP( 0, -1), cC = TAP( 1, -1);
    float3 eC = TAP(-1,  0), fC = TAP( 0,  0), gC = TAP( 1,  0), hC = TAP( 2,  0);
    float3 iC = TAP(-1,  1), jC = TAP( 0,  1), kC = TAP( 1,  1), lC = TAP( 2,  1);
    float3 nC = TAP( 0,  2), oC = TAP( 1,  2);
    #undef TAP

    float bL = edgeLuma(bC), cL = edgeLuma(cC);
    float eL = edgeLuma(eC), fL = edgeLuma(fC), gL = edgeLuma(gC), hL = edgeLuma(hC);
    float iL = edgeLuma(iC), jL = edgeLuma(jC), kL = edgeLuma(kC), lL = edgeLuma(lC);
    float nL = edgeLuma(nC), oL = edgeLuma(oC);

    // Four overlapping crosses, one per corner of the 2x2 {f,g,j,k}.
    float2 dir = float2(0.0f);
    float  len = 0.0f;
    easuSet(dir, len, pp, 0, bL, eL, fL, gL, jL);
    easuSet(dir, len, pp, 1, cL, fL, gL, hL, kL);
    easuSet(dir, len, pp, 2, fL, iL, jL, kL, nL);
    easuSet(dir, len, pp, 3, gL, jL, kL, lL, oL);

    // Normalise the direction. Below the threshold there is no meaningful edge,
    // so fall back to axis-aligned and let `len` (which will be ~0) neutralise
    // the anisotropy anyway.
    float2 dir2 = dir * dir;
    float  dirR = dir2.x + dir2.y;
    bool   zero = dirR < (1.0f / 32768.0f);
    dirR = zero ? 1.0f : rsqrt(dirR);
    dir.x = zero ? 1.0f : dir.x;
    dir  *= dirR;

    // len arrives in [0,2] from the four weighted crosses; remap to [0,1] and
    // square it so the kernel only reshapes on genuinely strong edges.
    len = len * 0.5f;
    len *= len;

    // Stretch the kernel from 1.0 on axis-aligned edges to sqrt(2) on diagonals,
    // so a 45-degree edge gets a filter footprint as wide as it deserves.
    float stretch = (dir.x * dir.x + dir.y * dir.y) * safeRcp(max(abs(dir.x), abs(dir.y)));

    // Anisotropic scaling: along the edge -> `stretch`, across the edge -> 0.5x
    // (tighter, so the edge stays crisp).
    float2 len2 = float2(1.0f + (stretch - 1.0f) * len,
                         1.0f - 0.5f * len);

    // Window size shrinks as the region becomes more edge-like.
    float lob = 0.5f + ((1.0f / 4.0f - 0.04f) - 0.5f) * len;
    float clp = safeRcp(lob);

    float3 accColor  = float3(0.0f);
    float  accWeight = 0.0f;
    easuTap(accColor, accWeight, float2( 0.0f, -1.0f) - pp, dir, len2, lob, clp, bC);
    easuTap(accColor, accWeight, float2( 1.0f, -1.0f) - pp, dir, len2, lob, clp, cC);
    easuTap(accColor, accWeight, float2(-1.0f,  1.0f) - pp, dir, len2, lob, clp, iC);
    easuTap(accColor, accWeight, float2( 0.0f,  1.0f) - pp, dir, len2, lob, clp, jC);
    easuTap(accColor, accWeight, float2( 0.0f,  0.0f) - pp, dir, len2, lob, clp, fC);
    easuTap(accColor, accWeight, float2(-1.0f,  0.0f) - pp, dir, len2, lob, clp, eC);
    easuTap(accColor, accWeight, float2( 1.0f,  1.0f) - pp, dir, len2, lob, clp, kC);
    easuTap(accColor, accWeight, float2( 2.0f,  1.0f) - pp, dir, len2, lob, clp, lC);
    easuTap(accColor, accWeight, float2( 2.0f,  0.0f) - pp, dir, len2, lob, clp, hC);
    easuTap(accColor, accWeight, float2( 1.0f,  0.0f) - pp, dir, len2, lob, clp, gC);
    easuTap(accColor, accWeight, float2( 1.0f,  2.0f) - pp, dir, len2, lob, clp, oC);
    easuTap(accColor, accWeight, float2( 0.0f,  2.0f) - pp, dir, len2, lob, clp, nC);

    // Deringing: the Lanczos lobes go negative, which can overshoot past the
    // brightest/darkest neighbour and produce halos. Clamping to the range of
    // the immediate 2x2 removes that without softening the edge itself.
    float3 mn4 = min(min(fC, gC), min(jC, kC));
    float3 mx4 = max(max(fC, gC), max(jC, kC));
    float3 pix = clamp(accColor * safeRcp(accWeight), mn4, mx4);

    dst.write(float4(pix, 1.0f), gid);
}

//===----------------------------------------------------------------------===//
//  Pass 2 — RCAS
//===----------------------------------------------------------------------===//

// RCAS sharpens by mixing in a ring of 4 neighbours with a *negative* weight
// `lobe`. The magnitude of that weight is capped per-pixel so the result
// provably cannot exceed the local min/max — that is the "robust" part, and it
// is why RCAS does not ring the way an unconstrained unsharp mask does.
kernel void csr_rcas(texture2d<float, access::read>  src [[texture(0)]],
                      texture2d<float, access::write> dst [[texture(1)]],
                      constant SpatialUniforms &u            [[buffer(0)]],
                      uint2 gid                           [[thread_position_in_grid]]) {
    uint w = dst.get_width();
    uint h = dst.get_height();
    if (gid.x >= w || gid.y >= h) { return; }

    // Clamped neighbour fetch.
    //     b
    //   d e f
    //     h
    int2 p  = int2(gid);
    int2 hi = int2(int(w) - 1, int(h) - 1);
    #define RD(dx, dy) src.read(uint2(clamp(p + int2(dx, dy), int2(0), hi))).rgb

    float3 b = RD( 0, -1);
    float3 d = RD(-1,  0);
    float3 e = RD( 0,  0);
    float3 f = RD( 1,  0);
    float3 hC= RD( 0,  1);
    #undef RD

    // RCAS denoising works on the GREEN channel alone, not the luma proxy EASU
    // uses. That is what the reference algorithm does, and it matters:
    // green carries most of the perceived detail, and using it keeps the noise
    // estimate independent of the red/blue chroma noise that would otherwise
    // make the detector suppress genuine sharpening on saturated content.
    float bL = b.g, dL = d.g, eL = e.g;
    float fL = f.g, hL = hC.g;

    // Noise detection: how far the centre sits from the mean of its ring,
    // relative to the ring's own spread. A lone bright pixel in a flat area
    // scores high and gets its sharpening halved, so film grain and dithering
    // do not get amplified.
    float nz = 0.25f * (bL + dL + fL + hL) - eL;
    float range = max3f(max3f(bL, dL, eL), fL, hL) - min3f(min3f(bL, dL, eL), fL, hL);
    nz = saturate(abs(nz) * safeRcp(range));
    nz = 1.0f - 0.5f * nz;

    // Per-channel ring extremes (centre excluded).
    float3 mn4 = min(min(b, d), min(f, hC));
    float3 mx4 = max(max(b, d), max(f, hC));

    // Solve, per channel, for the largest |lobe| that still keeps the sharpened
    // result inside [0,1] given this ring. Take the most restrictive channel.
    // Headroom against black: 4*mx4 is >= 0, so clamp away the zero case.
    float3 hitMin = mn4 / max(4.0f * mx4, float3(1e-6f));
    // Headroom against white: (4*mn4 - 4) is <= 0, so clamp it from above.
    float3 hitMax = (float3(1.0f) - mx4) / min(4.0f * mn4 - 4.0f, float3(-1e-6f));
    float3 lobeRGB = max(-hitMin, hitMax);

    // Hard sharpening limit: never let the lobe get stronger than this regardless of
    // what the per-channel solve permits.
    const float kRcasLimit = 0.25f - (1.0f / 16.0f);
    float lobe = max(-kRcasLimit,
                     min(max3f(lobeRGB.r, lobeRGB.g, lobeRGB.b), 0.0f));

    // `sharpness` is a stops value: 0 = full strength, higher = softer.
    lobe *= exp2(-u.sharpness);
    lobe *= nz;

    float rcpL = safeRcp(4.0f * lobe + 1.0f);
    float3 outC = (lobe * (b + d + f + hC) + e) * rcpL;

    dst.write(float4(saturate(outC), 1.0f), gid);
}
