//===----------------------------------------------------------------------===//
//  Temporal.metal - temporal accumulation upscaling.
//
//  WHAT THIS IS, HONESTLY
//  ---------------------
//  A renderer-integrated temporal upscaler is a different animal. It
//  requires three things from the application that drew the frame:
//
//      1. per-pixel motion vectors,
//      2. a depth buffer (for disocclusion detection and depth-aware clamping),
//      3. a known sub-pixel jitter sequence applied to the projection matrix.
//
//  We are upscaling an already-composited desktop captured by ScreenCaptureKit.
//  None of those three inputs exist, and none can be recovered faithfully — the
//  compositor has thrown that information away by the time we see pixels.
//
//  So this cannot be a renderer-integrated temporal upscaler, and is not one.
//  It keeps the same *structure* with the missing inputs derived instead of
//  supplied:
//
//      - motion vectors  -> estimated by the optical-flow pyramid in
//                           FrameGen.metal (coarse-to-fine block matching on
//                           luma). The same field that drives frame
//                           interpolation optical-flow stage uses.
//      - depth           -> unavailable. Replaced by colour-based disocclusion
//                           detection: where reprojected history disagrees
//                           strongly with the current frame, history is
//                           rejected.
//      - jitter          -> unavailable, and unfakeable. Without jitter there
//                           is no new sub-pixel information arriving each
//                           frame, so this cannot reconstruct genuinely unseen
//                           detail the way a jittered temporal upscaler does. What it does
//                           deliver is temporal stability: accumulated history
//                           removes the shimmer and crawling that a purely
//                           spatial upscaler shows on text and thin UI lines,
//                           which on desktop content is the dominant artefact.
//
//  PIPELINE POSITION
//  -----------------
//      current frame --[EASU, from Spatial.metal]--> spatial estimate
//                                                       |
//      history (previous output) --[reproject by motion]+--> accumulate
//                                                       |
//                                                  [RCAS] --> output
//
//  The accumulation pass below is the only new kernel; EASU and RCAS are shared
//  with the spatial path, which is why the two modes differ only in whether
//  variants relate.
//===----------------------------------------------------------------------===//

#include <metal_stdlib>
using namespace metal;

struct TemporalUniforms {
    float2 inputSize;      // captured source dimensions, pixels
    float2 outputSize;     // upscaled output dimensions, pixels
    float2 motionScale;    // motion-field texel -> output pixel conversion
    float  feedback;       // steady-state history weight, 0..1
    float  disoccl;        // disocclusion sensitivity; higher rejects sooner
    float  historyValid;   // 0 on the first frame after a reset, else 1
    float  _pad0;
};

//===----------------------------------------------------------------------===//
//  Neighbourhood rectification.
//
//  The standard temporal-AA safeguard, sometimes called the "rectified"
//  history. Reprojected history is only trustworthy if it looks like the
//  current frame's local neighbourhood. We build the axis-aligned colour
//  bounding box of the 3x3 neighbourhood around the current pixel and clamp
//  history into it. Anything that drifted outside — a stale ghost trailing a
//  moving window, say — gets pulled back onto the current frame's colour range
//  instead of smearing.
//===----------------------------------------------------------------------===//
struct ColorBox {
    float3 lo;
    float3 hi;
    float3 avg;
};

static inline ColorBox neighbourhoodBox(texture2d<float, access::sample> spatial,
                                        float2 uv,
                                        float2 invOutSize) {
    constexpr sampler s(coord::normalized, filter::linear, address::clamp_to_edge);

    ColorBox box;
    box.lo  = float3(1.0);
    box.hi  = float3(0.0);
    box.avg = float3(0.0);

    for (int y = -1; y <= 1; ++y) {
        for (int x = -1; x <= 1; ++x) {
            float3 c = spatial.sample(s, uv + float2(x, y) * invOutSize).rgb;
            box.lo   = min(box.lo, c);
            box.hi   = max(box.hi, c);
            box.avg += c;
        }
    }
    box.avg *= (1.0 / 9.0);
    return box;
}

//===----------------------------------------------------------------------===//
//  csr_accumulate — reproject history, rectify it, blend with the spatial
//  estimate.
//
//  texture(0)  spatial   : EASU output for the current frame, output-res
//  texture(1)  history   : accumulated output from the previous frame, output-res
//  texture(2)  motion    : optical-flow field, in *source* pixels per frame
//  texture(3)  dst       : accumulated result, output-res
//===----------------------------------------------------------------------===//
kernel void csr_accumulate(texture2d<float, access::sample> spatial [[texture(0)]],
                            texture2d<float, access::sample> history [[texture(1)]],
                            texture2d<float, access::sample> motion  [[texture(2)]],
                            texture2d<float, access::write>  dst     [[texture(3)]],
                            constant TemporalUniforms &u                 [[buffer(0)]],
                            uint2 gid [[thread_position_in_grid]]) {

    if (gid.x >= uint(u.outputSize.x) || gid.y >= uint(u.outputSize.y)) { return; }

    constexpr sampler linearClamp(coord::normalized, filter::linear, address::clamp_to_edge);

    float2 invOutSize = 1.0 / u.outputSize;
    float2 uv = (float2(gid) + 0.5) * invOutSize;

    float3 current = spatial.sample(linearClamp, uv).rgb;

    // First frame after a reset: nothing to accumulate against.
    if (u.historyValid < 0.5) {
        dst.write(float4(current, 1.0), gid);
        return;
    }

    //--- Reprojection -------------------------------------------------------
    // The motion field says, for this location, how far the content moved
    // between the previous and current capture, measured in source pixels.
    // Stepping *backwards* along it lands on where this pixel's content was in
    // the previous output — i.e. where its history lives.
    float2 mv = motion.sample(linearClamp, uv).xy;
    float2 historyUV = uv - mv * u.motionScale * invOutSize;

    // Off-screen history never existed, so it cannot be reused.
    bool offscreen = any(historyUV < 0.0) || any(historyUV > 1.0);

    float3 hist = history.sample(linearClamp, historyUV).rgb;

    //--- Rectification ------------------------------------------------------
    ColorBox box = neighbourhoodBox(spatial, uv, invOutSize);

    // Widen the box slightly. A box built from exactly 9 samples is tight
    // enough that even correct history gets clipped on noisy content, which
    // throws away the stability we are here to gain.
    float3 slack = (box.hi - box.lo) * 0.25 + 1.0e-3;
    float3 lo = box.lo - slack;
    float3 hi = box.hi + slack;

    float3 clampedHist = clamp(hist, lo, hi);

    //--- Disocclusion / confidence -----------------------------------------
    // How far history had to be moved to fit tells us how wrong it was. With
    // no depth buffer this colour-space distance is our only disocclusion
    // signal: newly revealed content (a window moving off something behind it)
    // has history that disagrees violently and must be dropped.
    float rejection = length(clampedHist - hist);
    float confidence = exp2(-u.disoccl * rejection);

    // Fast motion also lowers confidence: block-matched motion is least
    // reliable exactly where things move quickest, and stale history there
    // reads as a trail.
    float speed = length(mv);
    confidence *= exp2(-0.25 * speed);

    if (offscreen) { confidence = 0.0; }

    //--- Accumulate ---------------------------------------------------------
    // `feedback` is the steady-state history weight for a perfectly confident
    // pixel. Scaling it by confidence collapses the blend back to the purely
    // spatial estimate wherever history cannot be trusted, so the worst case
    // degrades to the spatial result, never to a smear.
    float alpha = u.feedback * confidence;
    float3 result = mix(current, clampedHist, alpha);

    dst.write(float4(result, 1.0), gid);
}

//===----------------------------------------------------------------------===//
//  csr_copy — plain output-res copy, used to seed and roll the history chain.
//===----------------------------------------------------------------------===//
kernel void csr_copy(texture2d<float, access::read>  src [[texture(0)]],
                      texture2d<float, access::write> dst [[texture(1)]],
                      uint2 gid [[thread_position_in_grid]]) {
    if (gid.x >= dst.get_width() || gid.y >= dst.get_height()) { return; }
    dst.write(src.read(gid), gid);
}

//===----------------------------------------------------------------------===//
//  csr_accumulate_hq — the high-quality accumulation path.
//
//  Same job as `csr_accumulate`, but with three upgrades that between them
//  raise the ceiling on how much history can safely be kept. The standard path
//  is conservative because its safeguard is crude; make the safeguard better
//  and you can accumulate far more aggressively without ghosting.
//
//    1. VARIANCE CLIPPING IN YCoCg, instead of an RGB min/max box.
//
//       A min/max box over 9 samples is the *convex hull* of the neighbourhood
//       — one outlier pixel inflates it and lets stale history through. Instead
//       we take the mean and standard deviation of the neighbourhood and build
//       the box as mean ± γ·σ, which describes where the samples actually
//       cluster rather than where the two extremes happen to fall.
//
//       Doing it in YCoCg rather than RGB matters: YCoCg separates luma from
//       two chroma axes, so the box is aligned with the direction the human eye
//       actually resolves detail in. An RGB box is three correlated axes, which
//       makes it simultaneously too tight on chroma and too loose on luma.
//
//    2. CLIPPING TOWARD THE CURRENT COLOUR, not per-channel clamping.
//
//       Clamping each channel independently moves the history colour to a
//       corner of the box and shifts its hue. Clipping instead walks the
//       straight line from the current colour toward the history colour and
//       stops at the box boundary, so the result keeps history's hue and only
//       loses its magnitude. This is what stops the faint colour fringing that
//       per-channel clamping leaves on moving edges.
//
//    3. ANTI-FLICKER LUMINANCE WEIGHTING.
//
//       Blending in linear-ish luma lets one very bright pixel dominate its
//       neighbours and flicker between frames. Weighting each contribution by
//       1/(1+Y) before the blend, and dividing it back out afterwards,
//       compresses that dominance — the standard tone-mapped-blend trick.
//===----------------------------------------------------------------------===//

static inline float3 rgbToYCoCg(float3 c) {
    return float3( 0.25f * c.r + 0.5f * c.g + 0.25f * c.b,
                   0.5f  * c.r              - 0.5f  * c.b,
                  -0.25f * c.r + 0.5f * c.g - 0.25f * c.b);
}

static inline float3 yCoCgToRGB(float3 c) {
    return float3(c.x + c.y - c.z,
                  c.x        + c.z,
                  c.x - c.y - c.z);
}

/// Walk from `anchor` toward `target` and stop at the AABB boundary.
/// Returns `target` untouched when it is already inside the box.
static inline float3 clipToBox(float3 anchor, float3 target,
                               float3 boxMin, float3 boxMax) {
    float3 centre = 0.5f * (boxMax + boxMin);
    float3 extent = 0.5f * (boxMax - boxMin) + 1.0e-6f;

    float3 offset = target - centre;
    // How many "extents" the target sits along each axis. The largest of the
    // three decides where the ray leaves the box.
    float3 units = abs(offset / extent);
    float  maxUnit = max(units.x, max(units.y, units.z));

    if (maxUnit > 1.0f) {
        return centre + offset / maxUnit;
    }
    return target;
}

kernel void csr_accumulate_hq(texture2d<float, access::sample> spatial [[texture(0)]],
                              texture2d<float, access::sample> history [[texture(1)]],
                              texture2d<float, access::sample> motion  [[texture(2)]],
                              texture2d<float, access::write>  dst     [[texture(3)]],
                              constant TemporalUniforms &u             [[buffer(0)]],
                              uint2 gid [[thread_position_in_grid]]) {

    if (gid.x >= uint(u.outputSize.x) || gid.y >= uint(u.outputSize.y)) { return; }

    constexpr sampler linearClamp(coord::normalized, filter::linear, address::clamp_to_edge);

    float2 invOutSize = 1.0 / u.outputSize;
    float2 uv = (float2(gid) + 0.5f) * invOutSize;

    float3 currentRGB = spatial.sample(linearClamp, uv).rgb;

    if (u.historyValid < 0.5f) {
        dst.write(float4(currentRGB, 1.0f), gid);
        return;
    }

    //--- Neighbourhood statistics in YCoCg ---------------------------------
    // First and second moments give mean and variance in one pass, which is
    // cheaper than a second loop and numerically fine at 8-bit input depth.
    float3 m1 = 0.0f;
    float3 m2 = 0.0f;
    for (int y = -1; y <= 1; ++y) {
        for (int x = -1; x <= 1; ++x) {
            float3 s = rgbToYCoCg(spatial.sample(linearClamp,
                                                 uv + float2(x, y) * invOutSize).rgb);
            m1 += s;
            m2 += s * s;
        }
    }
    const float inv9 = 1.0f / 9.0f;
    float3 mean  = m1 * inv9;
    float3 sigma = sqrt(max(m2 * inv9 - mean * mean, 0.0f));

    // γ = 1.25 keeps roughly 90% of a normal distribution. Tighter ghosts less
    // but reintroduces shimmer; looser is the reverse. This is the knob that
    // trades the two against each other.
    const float gamma = 1.25f;
    float3 boxMin = mean - gamma * sigma;
    float3 boxMax = mean + gamma * sigma;

    //--- Reprojection -------------------------------------------------------
    float2 mv = motion.sample(linearClamp, uv).xy;
    float2 historyUV = uv - mv * u.motionScale * invOutSize;
    bool offscreen = any(historyUV < 0.0f) || any(historyUV > 1.0f);

    float3 historyYCoCg = rgbToYCoCg(history.sample(linearClamp, historyUV).rgb);
    float3 currentYCoCg = rgbToYCoCg(currentRGB);

    // Clip along the line from current to history, preserving hue.
    float3 clipped = clipToBox(currentYCoCg, historyYCoCg, boxMin, boxMax);

    //--- Confidence ---------------------------------------------------------
    // Distance travelled during clipping, expressed in units of the local
    // standard deviation, so the measure is scale-free: a given displacement
    // means much more in a flat region than in a busy one.
    float3 delta = (clipped - historyYCoCg) / max(sigma, float3(1.0e-3f));
    float  rejection = length(delta);
    float  confidence = exp2(-u.disoccl * 0.35f * rejection);

    float speed = length(mv);
    confidence *= exp2(-0.25f * speed);
    if (offscreen) { confidence = 0.0f; }

    //--- Anti-flicker weighted blend ---------------------------------------
    // Weight by 1/(1+Y) so a very bright sample cannot dominate the blend, then
    // divide the weighting back out.
    float alpha = u.feedback * confidence;

    float wCur  = 1.0f / (1.0f + max(currentYCoCg.x, 0.0f));
    float wHist = 1.0f / (1.0f + max(clipped.x, 0.0f));

    float3 blended = (currentYCoCg * wCur * (1.0f - alpha) + clipped * wHist * alpha)
                   / max(wCur * (1.0f - alpha) + wHist * alpha, 1.0e-5f);

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