// Black/white normalisation and X-Trans demosaic, in one pass.
//
// The Fujifilm counterpart to demosaic.wgsl (FR-RAW-5). The input and the
// output are the same — packed u16 photosites in, linear camera-space
// RGBA16Float out — but the colour filter array is a 6x6 tile rather than a
// 2x2 one, and none of the Bayer kernels survive that. In a Bayer cell every
// pixel has the missing channels at a fixed offset; in X-Trans the offsets
// differ at all 36 positions, so a fixed kernel per site would need 36 of
// them and would still say nothing about which neighbours to trust.
//
// The method here is *local plane fitting on the colour difference*:
//
//  1. Take a 5x5 window and sort its photosites by the channel each one
//     measures. Every window holds at least four red, four blue and thirteen
//     green samples, whatever the phase — checked exhaustively, not assumed.
//  2. Fit a weighted least-squares plane through each channel's samples and
//     evaluate all three planes at the pixel. A plane rather than a mean
//     because the three channels are sampled at *different* places: a mean
//     would compare a red average taken slightly left of the pixel with a
//     green average taken slightly right of it, and the difference of those
//     two offsets is a colour cast that follows every gradient in the frame.
//     A plane has no such bias — it reconstructs any linear gradient exactly.
//  3. Keep the pixel's own measured value, and carry the other two channels
//     across as the *difference* between the fitted planes.
//
// Step 3 is the standard constant-colour-difference model, and the reason it
// is applied in white-balanced space is that the model is exact only where
// the channel difference is locally constant. On a neutral subject that is
// true after white balance and false before it, so the gains go on before the
// fit and come off after — which costs one multiply and halves the error at a
// luminance edge (measured on a synthetic step: 0.34 -> 0.20 max error).
// The pixels written out are still as-shot, unbalanced camera space; nothing
// downstream sees the difference.
//
// **What this is not.** It is not Markesteijn. It has no directional
// hypotheses and no homogeneity map, so it does not resolve detail finer than
// the CFA period, and a hard edge arrives about two pixels wide. It does not
// produce the "worms" that FR-RAW-5 exists to avoid — the output is bounded
// by the local sample range, so it cannot ring — but the Markesteijn-class
// quality that requirement asks for is still owed.

struct XTransParams {
    // Dimensions of the *cropped* output, in pixels.
    width: u32,
    height: u32,
    // Origin of the crop within the sensor readout, in photosites. Added to
    // every read, and to every pattern lookup: the 6x6 tile is anchored to the
    // sensor, not to the visible frame.
    crop_x: u32,
    crop_y: u32,
    // Row stride of the input, in samples.
    stride: u32,
    // One black level and one reciprocal range for the whole sensor. The
    // four-value form the Bayer path uses is a 2x2 convention with no meaning
    // on a 6x6 tile.
    black: f32,
    inv_range: f32,
    _pad0: u32,
    // As-shot white balance gains, green-normalised, and their reciprocals.
    wb: vec4<f32>,
    inv_wb: vec4<f32>,
    // The 6x6 tile, already rotated to this sensor's phase on the CPU, packed
    // two bits per photosite: word k holds row 2k in its low 12 bits and row
    // 2k+1 in the next 12. The fourth word is padding.
    tile: vec4<u32>,
}

@group(0) @binding(0) var<storage, read> raw: array<u32>;
@group(0) @binding(1) var<uniform> params: XTransParams;
@group(0) @binding(2) var output: texture_storage_2d<rgba16float, write>;

// Half-width of the fitting window. Two is the smallest radius for which
// every phase of the tile still offers enough red and blue samples to pin a
// plane down; three would be smoother and blurrier.
const RADIUS: i32 = 2;

// Colour of the photosite at absolute sensor coordinates: 0=R, 1=G, 2=B.
fn colour_at(sx: u32, sy: u32) -> u32 {
    let row = sy % 6u;
    let col = sx % 6u;
    let word = params.tile[row >> 1u];
    return (word >> ((row & 1u) * 12u + col * 2u)) & 3u;
}

// Read one photosite, normalised to [0, 1] against the black level.
//
// Coordinates are relative to the crop origin and must already be in range;
// unlike the Bayer pass there is no reflection here, because the window is
// slid inside the image instead (see `main`) and so never asks for a
// photosite that does not exist.
fn sample(cx: i32, cy: i32) -> f32 {
    let sx = u32(cx) + params.crop_x;
    let sy = u32(cy) + params.crop_y;
    let index = sy * params.stride + sx;

    let word = raw[index >> 1u];
    let raw_value = select(word & 0xFFFFu, word >> 16u, (index & 1u) == 1u);

    // Sensor noise puts real signal below the black point, so subtracting it
    // can go negative; clamped rather than allowed to wrap.
    return max((f32(raw_value) - params.black) * params.inv_range, 0.0);
}

// Solve the 3x3 weighted least-squares normal equations for a plane
// `c0 + c1*dx + c2*dy` and evaluate it at `(ex, ey)`.
//
// The accumulators are the usual moments: `n` is the summed weight, `sx`..`syy`
// the first and second moments of the sample positions, `t0`..`ty` the same
// moments weighted by value. Cramer's rule rather than a factorisation — the
// matrix is 3x3 and symmetric, and this keeps the whole solve in registers.
fn plane_at(
    n: f32, sx: f32, sy: f32, sxx: f32, sxy: f32, syy: f32,
    t0: f32, tx: f32, ty: f32, ex: f32, ey: f32,
) -> f32 {
    if (n <= 0.0) {
        return 0.0;
    }
    let det = n * (sxx * syy - sxy * sxy)
            - sx * (sx * syy - sxy * sy)
            + sy * (sx * sxy - sxx * sy);

    // Degenerate only if a channel's samples in this window are collinear,
    // which the 5x5 geometry rules out for every phase — but an image a few
    // photosites across is clipped down to fewer samples than that, and a
    // division by a near-zero determinant there would put NaN in the texture.
    // Falling back to the plain weighted mean loses the gradient term and
    // nothing else.
    if (abs(det) < 1e-6 * n * n * n) {
        return t0 / n;
    }

    let inv = 1.0 / det;
    let c0 = (t0 * (sxx * syy - sxy * sxy)
            - sx * (tx * syy - sxy * ty)
            + sy * (tx * sxy - sxx * ty)) * inv;
    let c1 = (n * (tx * syy - sxy * ty)
            - t0 * (sx * syy - sxy * sy)
            + sy * (sx * ty - tx * sy)) * inv;
    let c2 = (n * (sxx * ty - tx * sxy)
            - sx * (sx * ty - tx * sy)
            + t0 * (sx * sxy - sxx * sy)) * inv;
    return c0 + c1 * ex + c2 * ey;
}

@compute @workgroup_size(8, 8, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    if (gid.x >= params.width || gid.y >= params.height) {
        return;
    }

    let x = i32(gid.x);
    let y = i32(gid.y);
    let w = i32(params.width);
    let h = i32(params.height);

    // Near an edge the window slides inward rather than reflecting. Reflection
    // is right for the Bayer pass, whose kernels only need the mirrored sample
    // to be the same colour; here the fit needs real geometry, and a mirrored
    // photosite sitting at a position it does not occupy tilts the plane. A
    // slid window is entirely real data, so the plane stays exact right up to
    // the border — the pixel is still inside the window, just not centred in
    // it, which is what `ex`/`ey` below account for.
    let wx = clamp(x, RADIUS, max(RADIUS, w - 1 - RADIUS));
    let wy = clamp(y, RADIUS, max(RADIUS, h - 1 - RADIUS));

    // Per-channel moments, indexed 0=R, 1=G, 2=B.
    var n = vec3<f32>(0.0);
    var sx = vec3<f32>(0.0);
    var sy = vec3<f32>(0.0);
    var sxx = vec3<f32>(0.0);
    var sxy = vec3<f32>(0.0);
    var syy = vec3<f32>(0.0);
    var t0 = vec3<f32>(0.0);
    var tx = vec3<f32>(0.0);
    var ty = vec3<f32>(0.0);
    var lo = vec3<f32>(1.0e30);
    var hi = vec3<f32>(-1.0e30);

    for (var dy = -RADIUS; dy <= RADIUS; dy = dy + 1) {
        for (var dx = -RADIUS; dx <= RADIUS; dx = dx + 1) {
            // The clamp only bites on an image narrower than the window, where
            // a duplicated photosite is better than a missing channel.
            let px = clamp(wx + dx, 0, w - 1);
            let py = clamp(wy + dy, 0, h - 1);

            let v = sample(px, py);
            let k = colour_at(u32(px) + params.crop_x, u32(py) + params.crop_y);
            let u = v * params.wb[k];

            let fx = f32(px - wx);
            let fy = f32(py - wy);
            // Nearer photosites describe this pixel better. 1/(1+r^2) rather
            // than a Gaussian because it needs no width to tune and leaves the
            // normal equations well conditioned at every phase.
            let g = 1.0 / (1.0 + fx * fx + fy * fy);

            n[k] = n[k] + g;
            sx[k] = sx[k] + g * fx;
            sy[k] = sy[k] + g * fy;
            sxx[k] = sxx[k] + g * fx * fx;
            sxy[k] = sxy[k] + g * fx * fy;
            syy[k] = syy[k] + g * fy * fy;
            t0[k] = t0[k] + g * u;
            tx[k] = tx[k] + g * u * fx;
            ty[k] = ty[k] + g * u * fy;

            lo[k] = min(lo[k], v);
            hi[k] = max(hi[k], v);
        }
    }

    let ex = f32(x - wx);
    let ey = f32(y - wy);
    var p = vec3<f32>(
        plane_at(n.r, sx.r, sy.r, sxx.r, sxy.r, syy.r, t0.r, tx.r, ty.r, ex, ey),
        plane_at(n.g, sx.g, sy.g, sxx.g, sxy.g, syy.g, t0.g, tx.g, ty.g, ex, ey),
        plane_at(n.b, sx.b, sy.b, sxx.b, sxy.b, syy.b, t0.b, tx.b, ty.b, ex, ey),
    );

    // The pixel's own channel always has a sample — itself — so its plane is
    // always real, and it is the reference the other two are carried across
    // from.
    let centre = colour_at(u32(x) + params.crop_x, u32(y) + params.crop_y);
    let measured = sample(x, y);

    // A channel with no sample at all cannot happen in a 5x5 window; it can on
    // an image a few photosites across, where the window collapses. Such a
    // channel is given the reference plane, which renders the pixel grey
    // rather than arbitrary.
    let empty = n <= vec3<f32>(0.0);
    p = select(p, vec3<f32>(p[centre]), empty);
    lo = select(lo, vec3<f32>(0.0), empty);
    hi = select(hi, vec3<f32>(1.0e30), empty);

    // Constant colour difference, in white-balanced space, undone on the way
    // out. The measured channel comes back bit-for-bit: its own plane cancels.
    var rgb = (vec3<f32>(measured * params.wb[centre]) + p - vec3<f32>(p[centre]))
        * params.inv_wb.rgb;

    // Bound each channel by what was actually measured nearby. The plane
    // difference overshoots wherever the colour itself changes across the
    // window — a red edge against green — and an overshoot here is a coloured
    // halo. The pixel is always inside its own window, so a true value can
    // never be clipped away by this on smooth content.
    rgb = clamp(rgb, lo, hi);
    rgb = max(rgb, vec3<f32>(0.0));

    textureStore(output, vec2<i32>(x, y), vec4<f32>(rgb, 1.0));
}