//===----------------------------------------------------------------------===//
//  FrameGen.metal — motion-compensated frame interpolation ("frame boost").
//
//  Given two consecutive captured frames, synthesise one or more intermediate
//  frames so the output runs at 2x or 3x the capture rate.
//
//  APPROACH: hierarchical block matching, no machine learning. A coarse-to-fine
//  Gaussian pyramid of luma is searched to produce a dense-ish motion field,
//  which is then used to warp both source frames toward the intermediate time
//  and blend them. This is the same family of technique as classic broadcast
//  motion-compensated frame-rate conversion.
//
//  WHY NOT ML: the brief calls for a lightweight app with no heavy models. An
//  optical-flow network would dominate both bundle size and frame time. Block
//  matching runs in well under a millisecond at typical source resolutions and
//  needs no weights at all.
//
//  WHERE IT RUNS: on the *source* frames, before upscaling. Estimating motion
//  at 1440p and then upscaling the interpolated result is far cheaper than
//  upscaling twice, and the motion field is the same either way.
//
//  HONEST LIMITS: block matching has no notion of occlusion or transparency.
//  Disocclusions (background revealed behind a moving object) and rapidly
//  changing text will show brief artifacts. The consistency-based blend below
//  detects most of these and falls back toward a plain cross-dissolve, which
//  reads as a soft ghost rather than a torn edge. Frame generation also costs
//  one frame of latency by construction: you cannot interpolate toward a frame
//  you have not captured yet.
//===----------------------------------------------------------------------===//

#include <metal_stdlib>
using namespace metal;

constexpr sampler linClamp(coord::normalized, filter::linear,  address::clamp_to_edge);
constexpr sampler pntClamp(coord::normalized, filter::nearest, address::clamp_to_edge);

static inline float luma709(float3 c) {
    return dot(c, float3(0.2126f, 0.7152f, 0.0722f));
}

//===----------------------------------------------------------------------===//
//  Pyramid construction
//===----------------------------------------------------------------------===//

/// Colour -> luma, with a 2x box reduction folded in.
kernel void fg_luma_reduce(texture2d<float, access::sample> 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; }
    float2 uv = (float2(gid) + 0.5f) / float2(dst.get_width(), dst.get_height());
    // A single bilinear tap at the midpoint of the 2x2 already averages four
    // texels, so the reduction is one sample rather than four.
    dst.write(float4(luma709(src.sample(linClamp, uv).rgb)), gid);
}

/// Luma -> luma 2x reduction, for the deeper pyramid levels.
kernel void fg_luma_halve(texture2d<float, access::sample> 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; }
    float2 uv = (float2(gid) + 0.5f) / float2(dst.get_width(), dst.get_height());
    dst.write(float4(src.sample(linClamp, uv).r), gid);
}

//===----------------------------------------------------------------------===//
//  Motion estimation
//===----------------------------------------------------------------------===//

struct MotionParams {
    int   searchRadius;   // in texels of the current level
    float levelScale;     // multiply an incoming vector by this to enter this level (2.0)
    int   _pad0;
    int   _pad1;
};

/// Sum of absolute differences over a 5x5 patch, one frame against the other,
/// with `mv` (in level texels) applied symmetrically: prev is sampled backwards
/// along the vector and next forwards. Searching symmetrically about the centre
/// makes the resulting field naturally centred on the *interpolated* frame,
/// which is exactly where we need it.
static inline float patchSAD(texture2d<float, access::sample> prev,
                             texture2d<float, access::sample> next,
                             float2 uv, float2 mvUV) {
    float sad = 0.0f;
    for (int dy = -2; dy <= 2; ++dy) {
        for (int dx = -2; dx <= 2; ++dx) {
            float2 o = float2(dx, dy) * 1.0f;
            float2 texel = 1.0f / float2(prev.get_width(), prev.get_height());
            float2 base = uv + o * texel;
            float a = prev.sample(linClamp, base - mvUV * 0.5f).r;
            float b = next.sample(linClamp, base + mvUV * 0.5f).r;
            sad += abs(a - b);
        }
    }
    return sad;
}

/// Coarsest level: exhaustive search with no prior.
kernel void fg_motion_init(texture2d<float, access::sample> prev [[texture(0)]],
                           texture2d<float, access::sample> next [[texture(1)]],
                           texture2d<float, access::write>  mvOut[[texture(2)]],
                           constant MotionParams &p             [[buffer(0)]],
                           uint2 gid [[thread_position_in_grid]]) {
    uint W = mvOut.get_width(), H = mvOut.get_height();
    if (gid.x >= W || gid.y >= H) { return; }

    float2 size  = float2(W, H);
    float2 texel = 1.0f / size;
    float2 uv    = (float2(gid) + 0.5f) * texel;

    float2 best = float2(0.0f);
    float  bestCost = patchSAD(prev, next, uv, float2(0.0f));

    for (int dy = -p.searchRadius; dy <= p.searchRadius; ++dy) {
        for (int dx = -p.searchRadius; dx <= p.searchRadius; ++dx) {
            if (dx == 0 && dy == 0) { continue; }
            float2 mv = float2(dx, dy);
            float cost = patchSAD(prev, next, uv, mv * texel);
            // Mild penalty on vector magnitude: with two equally good matches,
            // prefer the smaller motion. Suppresses a lot of spurious flow in
            // flat or periodic regions (scrollbars, window chrome, gradients).
            cost += 0.002f * length(mv);
            if (cost < bestCost) { bestCost = cost; best = mv; }
        }
    }
    // Store in normalised-UV units so the field can be sampled at any
    // resolution downstream without rescaling.
    mvOut.write(float4(best * texel, 0.0f, 0.0f), gid);
}

/// Finer levels: seed from the level above, then search a small window around
/// it. Coarse-to-fine is what lets a +-2 refine window still track large motion.
kernel void fg_motion_refine(texture2d<float, access::sample> prev  [[texture(0)]],
                             texture2d<float, access::sample> next  [[texture(1)]],
                             texture2d<float, access::sample> mvIn  [[texture(2)]],
                             texture2d<float, access::write>  mvOut [[texture(3)]],
                             constant MotionParams &p              [[buffer(0)]],
                             uint2 gid [[thread_position_in_grid]]) {
    uint W = mvOut.get_width(), H = mvOut.get_height();
    if (gid.x >= W || gid.y >= H) { return; }

    float2 size  = float2(W, H);
    float2 texel = 1.0f / size;
    float2 uv    = (float2(gid) + 0.5f) * texel;

    // The incoming field is already in UV units, so it carries across levels
    // unchanged — only the search step gets finer.
    float2 seed = mvIn.sample(linClamp, uv).xy;

    float2 best     = seed;
    float  bestCost = patchSAD(prev, next, uv, seed);

    // Also consider "no motion" — lets the field recover when the coarse level
    // locked onto a bad match.
    float zeroCost = patchSAD(prev, next, uv, float2(0.0f));
    if (zeroCost < bestCost) { bestCost = zeroCost; best = float2(0.0f); }

    for (int dy = -p.searchRadius; dy <= p.searchRadius; ++dy) {
        for (int dx = -p.searchRadius; dx <= p.searchRadius; ++dx) {
            if (dx == 0 && dy == 0) { continue; }
            float2 mv = seed + float2(dx, dy) * texel;
            float cost = patchSAD(prev, next, uv, mv);
            if (cost < bestCost) { bestCost = cost; best = mv; }
        }
    }
    mvOut.write(float4(best, 0.0f, 0.0f), gid);
}

/// 3x3 vector smoothing. Block matching produces a blocky field; without this
/// the interpolated frame shows visible block seams along moving edges.
kernel void fg_motion_smooth(texture2d<float, access::sample> mvIn  [[texture(0)]],
                             texture2d<float, access::write>  mvOut [[texture(1)]],
                             uint2 gid [[thread_position_in_grid]]) {
    uint W = mvOut.get_width(), H = mvOut.get_height();
    if (gid.x >= W || gid.y >= H) { return; }
    float2 texel = 1.0f / float2(W, H);
    float2 uv = (float2(gid) + 0.5f) * texel;

    float2 sum = float2(0.0f);
    float  wsum = 0.0f;
    for (int dy = -1; dy <= 1; ++dy) {
        for (int dx = -1; dx <= 1; ++dx) {
            float w = (dx == 0 && dy == 0) ? 4.0f : ((dx == 0 || dy == 0) ? 2.0f : 1.0f);
            sum  += mvIn.sample(pntClamp, uv + float2(dx, dy) * texel).xy * w;
            wsum += w;
        }
    }
    mvOut.write(float4(sum / wsum, 0.0f, 0.0f), gid);
}

//===----------------------------------------------------------------------===//
//  Interpolation
//===----------------------------------------------------------------------===//

struct InterpParams {
    float t;              // 0 = prev, 1 = next
    float sensitivity;    // 0..1, how eagerly to fall back to cross-dissolve
    float _pad0;
    float _pad1;
};

kernel void fg_interpolate(texture2d<float, access::sample> prev [[texture(0)]],
                           texture2d<float, access::sample> next [[texture(1)]],
                           texture2d<float, access::sample> mv   [[texture(2)]],
                           texture2d<float, access::write>  dst  [[texture(3)]],
                           constant InterpParams &p             [[buffer(0)]],
                           uint2 gid [[thread_position_in_grid]]) {
    uint W = dst.get_width(), H = dst.get_height();
    if (gid.x >= W || gid.y >= H) { return; }

    float2 uv = (float2(gid) + 0.5f) / float2(W, H);
    float2 v  = mv.sample(linClamp, uv).xy;   // prev -> next, in UV units

    // Warp both ways toward time t. Each source is pulled the fraction of the
    // vector that corresponds to its temporal distance from t.
    float3 a = prev.sample(linClamp, uv - v * p.t).rgb;
    float3 b = next.sample(linClamp, uv + v * (1.0f - p.t)).rgb;

    // Occlusion / bad-match detection. If the two warped samples disagree, the
    // motion vector did not explain this pixel: something was revealed,
    // covered, or simply changed (a cursor blink, a video cut). Blend toward a
    // straight cross-dissolve there, which degrades to a soft ghost instead of
    // a torn or smeared edge.
    float err = length(a - b);
    float conf = saturate(1.0f - err * mix(2.0f, 8.0f, p.sensitivity));

    float3 warped = mix(a, b, p.t);
    float3 plain  = mix(prev.sample(linClamp, uv).rgb,
                        next.sample(linClamp, uv).rgb, p.t);

    dst.write(float4(mix(plain, warped, conf), 1.0f), gid);
}
