// Watershed segmentation — the passes behind arm A of S15 (docs/dev/segmentation.md).
//
// Seven entry points forming one chain:
//
//   features      source texture -> perceptual triple, downscaled to proxy size
//   blur          pre-smoothing, without which every grain becomes a basin
//   gradient      Sobel magnitude — the surface the watershed floods
//   plateau_init  seed the distance field at every real descent
//   plateau_step  carry it inward, so flat ground drains toward its exit
//   flow          each pixel points downhill to its steepest neighbour
//   jump          pointer-jumping, until every pixel points at its basin root
//
// Everything after `features` works in storage buffers rather than textures.
// That is deliberate: the flow and jump passes need read-write access to the
// same array across dispatches, which storage textures do not give portably,
// and a buffer reads back without the 256-byte row padding a texture copy
// imposes.

struct Params {
    // Proxy dimensions — what every pass but `features` iterates over.
    width: u32,
    height: u32,
    // Source dimensions, for the box downscale in `features`.
    src_width: u32,
    src_height: u32,
    // Half-width of the pre-smoothing kernel, in proxy pixels. 0 disables it.
    blur_radius: i32,
    // 1 when the source is already display-encoded (the JPEG path), 0 for
    // linear scene-referred data out of the demosaicer.
    non_linear: u32,
    // How much luma and chroma each contribute to the gradient. Chroma is
    // weighted lower because it carries most of the sensor noise and few of
    // the boundaries a person would draw.
    w_luma: f32,
    w_chroma: f32,
}

// Binding slots are unique across the whole module, not reused per entry
// point: WGSL resource variables share one namespace, so two globals at the
// same (group, binding) is a module-level validation error even when no
// single entry point uses both. Each pass therefore gets its own pair, and
// each pipeline a layout declaring only the slots it touches.
@group(0) @binding(0) var<uniform> u: Params;

// ---------------------------------------------------------------- features

@group(0) @binding(1) var src: texture_2d<f32>;
@group(0) @binding(2) var<storage, read_write> feat_out: array<vec4<f32>>;

// Linear or display-encoded RGB to a roughly perceptual opponent triple.
//
// Perceptual rather than linear because the gradient has to agree with what
// a person calls an edge. In linear light a highlight rolloff swamps the
// boundary between two midtones, and the watershed would put its strongest
// walls where nobody sees one.
//
// The two chroma axes are opponent differences rather than a real Lab
// transform: they cost three subtractions instead of a matrix and a cube
// root, and the watershed only needs the *magnitude* of colour change, not a
// colorimetrically defensible value for it.
fn perceptual(c_in: vec3<f32>) -> vec3<f32> {
    var c = max(c_in, vec3<f32>(0.0));
    if (u.non_linear == 0u) {
        c = pow(c, vec3<f32>(1.0 / 2.4));
    }
    let l = dot(c, vec3<f32>(0.2126, 0.7152, 0.0722));
    let a = c.r - c.g;
    let b = c.b - 0.5 * (c.r + c.g);
    return vec3<f32>(l, a, b);
}

// Source -> proxy, averaging every source pixel that falls in the proxy
// pixel's footprint.
//
// A box average rather than point sampling because the proxy is where the
// segmentation happens: point sampling a 24 MP sensor down to 2 MP aliases
// fine texture into false gradient, and the watershed would faithfully find
// basins in the aliasing.
@compute @workgroup_size(8, 8, 1)
fn features(@builtin(global_invocation_id) gid: vec3<u32>) {
    if (gid.x >= u.width || gid.y >= u.height) {
        return;
    }

    let sx0 = (gid.x * u.src_width) / u.width;
    let sy0 = (gid.y * u.src_height) / u.height;
    let sx1 = max(sx0 + 1u, ((gid.x + 1u) * u.src_width) / u.width);
    let sy1 = max(sy0 + 1u, ((gid.y + 1u) * u.src_height) / u.height);

    var acc = vec3<f32>(0.0);
    var n = 0.0;
    for (var sy = sy0; sy < sy1; sy = sy + 1u) {
        for (var sx = sx0; sx < sx1; sx = sx + 1u) {
            let c = textureLoad(src, vec2<i32>(i32(sx), i32(sy)), 0).rgb;
            acc = acc + perceptual(c);
            n = n + 1.0;
        }
    }

    feat_out[gid.y * u.width + gid.x] = vec4<f32>(acc / max(n, 1.0), 0.0);
}

// -------------------------------------------------------------------- blur

@group(0) @binding(3) var<storage, read> blur_in: array<vec4<f32>>;
@group(0) @binding(4) var<storage, read_write> blur_out: array<vec4<f32>>;

fn clamp_coord(v: i32, hi: u32) -> u32 {
    return u32(clamp(v, 0, i32(hi) - 1));
}

// Pre-smoothing. Not a refinement — without it the watershed is unusable.
//
// A raw gradient over sensor data has a local minimum at every noise grain,
// and one basin per local minimum means a 2 MP frame segments into hundreds
// of thousands of regions that correspond to nothing. The radius is the
// caller's to set from ISO.
@compute @workgroup_size(8, 8, 1)
fn blur(@builtin(global_invocation_id) gid: vec3<u32>) {
    if (gid.x >= u.width || gid.y >= u.height) {
        return;
    }
    let idx = gid.y * u.width + gid.x;

    if (u.blur_radius <= 0) {
        blur_out[idx] = blur_in[idx];
        return;
    }

    var acc = vec3<f32>(0.0);
    var wsum = 0.0;
    let r = u.blur_radius;
    for (var dy = -r; dy <= r; dy = dy + 1) {
        for (var dx = -r; dx <= r; dx = dx + 1) {
            let sx = clamp_coord(i32(gid.x) + dx, u.width);
            let sy = clamp_coord(i32(gid.y) + dy, u.height);
            let d2 = f32(dx * dx + dy * dy);
            let w = exp(-d2 / (2.0 * f32(r) * f32(r)));
            acc = acc + blur_in[sy * u.width + sx].rgb * w;
            wsum = wsum + w;
        }
    }
    blur_out[idx] = vec4<f32>(acc / wsum, 0.0);
}

// ---------------------------------------------------------------- gradient

@group(0) @binding(5) var<storage, read> grad_in: array<vec4<f32>>;
@group(0) @binding(6) var<storage, read_write> grad_out: array<f32>;

fn feat_at(x: i32, y: i32) -> vec3<f32> {
    let sx = clamp_coord(x, u.width);
    let sy = clamp_coord(y, u.height);
    return grad_in[sy * u.width + sx].rgb;
}

// Sobel magnitude over the weighted opponent triple.
//
// This is the surface the watershed floods, so its units matter for nothing
// except ordering — only the *relative* height of one boundary against
// another decides which regions merge first.
@compute @workgroup_size(8, 8, 1)
fn gradient(@builtin(global_invocation_id) gid: vec3<u32>) {
    if (gid.x >= u.width || gid.y >= u.height) {
        return;
    }
    let x = i32(gid.x);
    let y = i32(gid.y);

    let tl = feat_at(x - 1, y - 1);
    let tc = feat_at(x, y - 1);
    let tr = feat_at(x + 1, y - 1);
    let ml = feat_at(x - 1, y);
    let mr = feat_at(x + 1, y);
    let bl = feat_at(x - 1, y + 1);
    let bc = feat_at(x, y + 1);
    let br = feat_at(x + 1, y + 1);

    let gx = (tr + 2.0 * mr + br) - (tl + 2.0 * ml + bl);
    let gy = (bl + 2.0 * bc + br) - (tl + 2.0 * tc + tr);

    let w = vec3<f32>(u.w_luma, u.w_chroma, u.w_chroma);
    let wx = gx * w;
    let wy = gy * w;

    grad_out[gid.y * u.width + gid.x] = sqrt(dot(wx, wx) + dot(wy, wy));
}

// ---------------------------------------------------------- lower-complete
//
// A watershed needs every non-minimum pixel to have a lower neighbour. A real
// gradient does not oblige: a flat wall, a clipped sky or the inside of a
// uniform object is a **plateau**, where every neighbour is exactly equal and
// there is no downhill direction to follow.
//
// Left alone, the tie-break in `flow` sends every plateau pixel to its
// lowest-indexed neighbour, which is up and to the left. Each pixel therefore
// walks diagonally until it falls off the plateau, and one flat region becomes
// a fan of diagonal chains rather than one basin — visible as hatching across
// what should be a single area (docs/dev/segmentation.md §12, F1).
//
// The fix is the standard lower-completion: give each plateau pixel its
// geodesic distance to the nearest pixel that *does* have a lower neighbour,
// then let `flow` order on (gradient, distance). Water on a plateau now runs
// toward the plateau's exit, which is what it would physically do.
//
// A plateau with no exit at all is a genuine regional minimum — the inside of
// a uniform disc, say. Those pixels keep `PLATEAU_UNRESOLVED`, tie with each
// other, and fall through to the index tie-break, which collapses the whole
// connected plateau onto its lowest-indexed pixel. One basin, which is the
// right answer for a regional minimum.

const PLATEAU_UNRESOLVED: u32 = 0xffffffffu;

// How close two gradients must be to count as the same level.
//
// **Exact equality does not work here, and that is not a rounding nicety.**
// The gradient is a float computed from 8-bit samples, so a region the eye
// and the algorithm both consider flat still has neighbours differing in the
// sixth decimal. With `==`, the breadth-first step never advances past its
// seeds and the whole pass is a no-op; with `<`, nearly every pixel finds
// some marginally lower neighbour and is seeded at zero, which is the same
// no-op wearing a different hat. Both were measured before this constant
// existed.
//
// Sized against the gradient's own scale: features are normalised to 0..1, so
// a Sobel magnitude runs to a few units, and 1e-4 is far below any step a
// real edge produces while sitting comfortably above f32 noise from a blur.
const LEVEL_EPS: f32 = 1e-4;

// Whether `b` lies below `a` by more than the level tolerance.
fn strictly_below(b: f32, a: f32) -> bool {
    return b < a - LEVEL_EPS;
}

// Whether two gradients belong to the same plateau.
fn same_level(a: f32, b: f32) -> bool {
    return abs(a - b) <= LEVEL_EPS;
}

@group(0) @binding(7) var<storage, read> pinit_grad: array<f32>;
@group(0) @binding(8) var<storage, read_write> pinit_out: array<u32>;

// Seed the distance field: zero where a real descent exists, unresolved on a
// plateau.
@compute @workgroup_size(8, 8, 1)
fn plateau_init(@builtin(global_invocation_id) gid: vec3<u32>) {
    if (gid.x >= u.width || gid.y >= u.height) {
        return;
    }
    let idx = gid.y * u.width + gid.x;
    let here = pinit_grad[idx];

    for (var dy = -1; dy <= 1; dy = dy + 1) {
        for (var dx = -1; dx <= 1; dx = dx + 1) {
            if (dx == 0 && dy == 0) {
                continue;
            }
            let nx = i32(gid.x) + dx;
            let ny = i32(gid.y) + dy;
            if (nx < 0 || ny < 0 || nx >= i32(u.width) || ny >= i32(u.height)) {
                continue;
            }
            if (strictly_below(pinit_grad[u32(ny) * u.width + u32(nx)], here)) {
                pinit_out[idx] = 0u;
                return;
            }
        }
    }

    pinit_out[idx] = PLATEAU_UNRESOLVED;
}

@group(0) @binding(9) var<storage, read> pstep_grad: array<f32>;
@group(0) @binding(10) var<storage, read> pstep_in: array<u32>;
@group(0) @binding(11) var<storage, read_write> pstep_out: array<u32>;

// One breadth-first step inward from the plateau's rim.
//
// Iterated by the host a fixed number of times rather than to convergence: a
// convergence test costs a readback per pass, and the count only has to cover
// the widest plateau in the frame. Pixels still unresolved when the budget
// runs out keep `PLATEAU_UNRESOLVED` and behave exactly as they did before
// this pass existed — the degradation is graceful, not a wrong answer.
@compute @workgroup_size(8, 8, 1)
fn plateau_step(@builtin(global_invocation_id) gid: vec3<u32>) {
    if (gid.x >= u.width || gid.y >= u.height) {
        return;
    }
    let idx = gid.y * u.width + gid.x;
    let current = pstep_in[idx];
    if (current != PLATEAU_UNRESOLVED) {
        pstep_out[idx] = current;
        return;
    }

    let here = pstep_grad[idx];
    var best = PLATEAU_UNRESOLVED;
    for (var dy = -1; dy <= 1; dy = dy + 1) {
        for (var dx = -1; dx <= 1; dx = dx + 1) {
            if (dx == 0 && dy == 0) {
                continue;
            }
            let nx = i32(gid.x) + dx;
            let ny = i32(gid.y) + dy;
            if (nx < 0 || ny < 0 || nx >= i32(u.width) || ny >= i32(u.height)) {
                continue;
            }
            let ni = u32(ny) * u.width + u32(nx);
            // Only within the same plateau: a neighbour at a different height
            // is across a boundary, and its distance says nothing about the
            // way out of this one.
            if (!same_level(pstep_grad[ni], here)) {
                continue;
            }
            let nd = pstep_in[ni];
            if (nd != PLATEAU_UNRESOLVED && nd < best) {
                best = nd;
            }
        }
    }

    if (best == PLATEAU_UNRESOLVED) {
        pstep_out[idx] = PLATEAU_UNRESOLVED;
    } else {
        pstep_out[idx] = best + 1u;
    }
}

// -------------------------------------------------------------------- flow

@group(0) @binding(12) var<storage, read> flow_grad: array<f32>;
@group(0) @binding(13) var<storage, read> flow_dist: array<u32>;
@group(0) @binding(14) var<storage, read_write> flow_out: array<u32>;

// Each pixel points at the steepest-descent neighbour among its 8, or at
// itself if it is a local minimum — a basin seed.
//
// **The ordering is load-bearing, twice over.** Comparing on (gradient,
// plateau distance, index) rather than gradient alone gives a strict total
// order, so the pointer graph descends monotonically and cannot contain a
// cycle — plateaux, which are everywhere in a smoothed image, would otherwise
// make two equal pixels point at each other and hang the pointer-jumping
// below.
//
// It is also what makes the result reproducible. S15's M5 asks whether a
// label field is stable enough across GPU vendors to be a cache key
// (ARCH §6.13); an arbitrary tie-break would answer no before the question
// was asked.
@compute @workgroup_size(8, 8, 1)
fn flow(@builtin(global_invocation_id) gid: vec3<u32>) {
    if (gid.x >= u.width || gid.y >= u.height) {
        return;
    }
    let idx = gid.y * u.width + gid.x;

    var best_val = flow_grad[idx];
    var best_dist = flow_dist[idx];
    var best_idx = idx;

    for (var dy = -1; dy <= 1; dy = dy + 1) {
        for (var dx = -1; dx <= 1; dx = dx + 1) {
            if (dx == 0 && dy == 0) {
                continue;
            }
            let nx = i32(gid.x) + dx;
            let ny = i32(gid.y) + dy;
            if (nx < 0 || ny < 0 || nx >= i32(u.width) || ny >= i32(u.height)) {
                continue;
            }
            let ni = u32(ny) * u.width + u32(nx);
            let nv = flow_grad[ni];
            let nd = flow_dist[ni];

            // Lexicographic on (gradient, plateau distance, index) rather
            // than one fused scalar. Folding the distance into the gradient
            // as a small epsilon would need a scale factor that is small
            // enough never to cross a real gradient step and large enough to
            // survive f32 — a tuning problem with a silent failure mode,
            // where three explicit keys have neither.
            // The same tolerance the plateau passes use, and for the same
            // reason: with exact equality this tie never fires on real data,
            // so the distance carried inward above would be computed and then
            // never consulted — the pass measurably did nothing.
            var better = false;
            if (strictly_below(nv, best_val)) {
                better = true;
            } else if (same_level(nv, best_val)) {
                if (nd < best_dist) {
                    better = true;
                } else if (nd == best_dist && ni < best_idx) {
                    better = true;
                }
            }

            if (better) {
                best_val = nv;
                best_dist = nd;
                best_idx = ni;
            }
        }
    }

    flow_out[idx] = best_idx;
}

// -------------------------------------------------------------------- jump

@group(0) @binding(15) var<storage, read> jump_in: array<u32>;
@group(0) @binding(16) var<storage, read_write> jump_out: array<u32>;

// Pointer jumping: parent = parent[parent].
//
// Halves every path length per dispatch, so ceil(log2(longest path)) passes
// resolve every pixel to its basin root. The host runs a fixed count bounded
// by log2(pixel count) rather than testing for convergence, because a
// convergence test costs a readback per iteration and the bound is ~21
// dispatches of a trivial kernel.
@compute @workgroup_size(8, 8, 1)
fn jump(@builtin(global_invocation_id) gid: vec3<u32>) {
    if (gid.x >= u.width || gid.y >= u.height) {
        return;
    }
    let idx = gid.y * u.width + gid.x;
    jump_out[idx] = jump_in[jump_in[idx]];
}