// Black/white normalisation and Bayer demosaic, in one pass.
//
// Input is the raw sensor readout as packed u16 samples — one per photosite,
// in CFA order. Output is linear scene-referred RGBA16Float in *camera*
// colour space; the camera→sRGB matrix belongs to the adjust pass, so this
// stage is purely about reconstructing three channels from one.
//
// The algorithm is Malvar-He-Cutler (ICASSP 2004): bilinear interpolation
// plus a Laplacian correction taken from the channel that *is* sampled at
// each site. One 5x5 neighbourhood per pixel, and dramatically better than
// bilinear on edges — bilinear leaves visible zippering on any high-contrast
// boundary, which on a 24 MP file is the first thing seen at 1:1.
//
// Kernel coefficients below are the paper's, all over 8.

struct DemosaicParams {
    // 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 so the masked border is never sampled.
    crop_x: u32,
    crop_y: u32,
    // Row stride of the input, in samples.
    stride: u32,
    // CFA layout of the *cropped* image, already re-phased for the crop
    // origin by dr-decode: 0=RGGB, 1=BGGR, 2=GRBG, 3=GBRG.
    pattern: u32,
    _pad0: u32,
    _pad1: u32,
    // Per-CFA-position black levels, indexed by (y&1)*2 + (x&1).
    black: vec4<f32>,
    // Reciprocal of (white - black) per position, precomputed on the CPU so
    // the shader does no division.
    inv_range: vec4<f32>,
}

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

// Colour of the photosite at (x, y): 0=R, 1=G, 2=B.
//
// Each pattern is its 2x2 cell read row-major, packed two bits per entry so
// the lookup is an index and a shift rather than a branch.
fn colour_at(x: u32, y: u32) -> u32 {
    let cell = (y & 1u) * 2u + (x & 1u);
    // RGGB = R,G,G,B -> 0,1,1,2 ; BGGR = 2,1,1,0 ; GRBG = 1,0,2,1 ; GBRG = 1,2,0,1
    // Entry i occupies bits [2i, 2i+1], so the cell order reads
    // right-to-left in hex. Verified against a table rather than derived by
    // eye — two of these were wrong on the first attempt.
    var packed: u32;
    switch params.pattern {
        case 0u: { packed = 0x94u; }  // RGGB -> [0,1,1,2]
        case 1u: { packed = 0x16u; }  // BGGR -> [2,1,1,0]
        case 2u: { packed = 0x61u; }  // GRBG -> [1,0,2,1]
        default: { packed = 0x49u; }  // GBRG -> [1,2,0,1]
    }
    return (packed >> (cell * 2u)) & 3u;
}

// Whether the row through (x, y) is one carrying red photosites.
//
// Needed at green sites, where red lies along one axis and blue along the
// other, and which is which depends on the pattern.
fn red_is_horizontal(x: u32, y: u32) -> bool {
    // The horizontal neighbour of a green site.
    return colour_at(x + 1u, y) == 0u;
}

// Read one photosite, normalised to [0, 1] against its own black level.
//
// Coordinates are relative to the crop origin. A 5x5 window at the image edge
// reflects rather than reading masked photosites or running off the buffer.
fn sample(ix: i32, iy: i32) -> f32 {
    let w = i32(params.width);
    let h = i32(params.height);

    // Reflect at the borders, preserving CFA parity: reflecting by an even
    // distance keeps the mirrored sample the same colour as the one it
    // stands in for. Clamping instead would flatten the correction term and
    // leave a visible one-pixel seam along each edge.
    var cx = ix;
    var cy = iy;
    if (cx < 0) { cx = -cx; }
    if (cy < 0) { cy = -cy; }
    if (cx > w - 1) { cx = 2 * (w - 1) - cx; }
    if (cy > h - 1) { cy = 2 * (h - 1) - cy; }
    cx = clamp(cx, 0, w - 1);
    cy = clamp(cy, 0, h - 1);

    let sx = u32(cx) + params.crop_x;
    let sy = u32(cy) + params.crop_y;
    let index = sy * params.stride + sx;

    // Samples are u16, packed two per u32 word.
    let word = raw[index >> 1u];
    let raw_value = select(word & 0xFFFFu, word >> 16u, (index & 1u) == 1u);

    // Black level and range are per CFA position. Subtracting black can go
    // negative on sensor noise — real signal below the black point — so the
    // result is clamped rather than allowed to wrap.
    let cell = (u32(cy) & 1u) * 2u + (u32(cx) & 1u);
    let value = (f32(raw_value) - params.black[cell]) * params.inv_range[cell];

    // **Clamped at the top as well, and that is what stops blown highlights
    // going pink.** Sensors read above their declared white level — on a
    // Canon 6D CR2 the data reaches 16383 against a white of 15070 — so a
    // saturated pixel normalises to about 1.1 rather than 1.0.
    //
    // Left unclamped it survives the white balance, where red is multiplied
    // by ~1.93 and blue by ~1.68 against green's 1.0, and then the camera
    // matrix. Red and blue clip at the end of the pipeline; green, whose
    // matrix row is far less positive-heavy, does not. Red and blue high with
    // green low is magenta, and a clipped highlight came back pink.
    //
    // Clamping here makes a blown pixel saturate *neutrally*: all three
    // channels reach 1.0 together and the highlight is white, which is what a
    // blown highlight looks like and what every other developer produces.
    return clamp(value, 0.0, 1.0);
}

@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 c = sample(x, y);

    // 5x5 neighbourhood.
    let n1 = sample(x, y - 1);
    let s1 = sample(x, y + 1);
    let w1 = sample(x - 1, y);
    let e1 = sample(x + 1, y);

    let n2 = sample(x, y - 2);
    let s2 = sample(x, y + 2);
    let w2 = sample(x - 2, y);
    let e2 = sample(x + 2, y);

    let nw = sample(x - 1, y - 1);
    let ne = sample(x + 1, y - 1);
    let sw = sample(x - 1, y + 1);
    let se = sample(x + 1, y + 1);

    let axial1 = n1 + s1 + w1 + e1;
    let diag1 = nw + ne + sw + se;
    let vert2 = n2 + s2;
    let horiz2 = w2 + e2;

    let colour = colour_at(gid.x, gid.y);
    var rgb: vec3<f32>;

    if (colour == 1u) {
        // ---- Green site ----------------------------------------------
        // Green is measured. Red and blue are interpolated from their own
        // axis, with a correction from the green Laplacian.
        //
        // Malvar "R at green in R row" kernel, and its transpose:
        //   chroma along the row:  (5c + 4(w1+e1) - (nw+ne+sw+se) - (w2+e2) + 0.5(n2+s2)) / 8
        //
        // The -1 goes on the two greens *along* the chroma axis and the +0.5
        // on the pair across it. Transposed, both kernels still sum to zero
        // and reconstruct a flat patch exactly, but on an edge the correction
        // at green sites is half strength and the false colour doubles: a
        // blue/yellow zipper around every clipped highlight.
        let along_row =
            (5.0 * c + 4.0 * (w1 + e1) - diag1 - horiz2 + 0.5 * vert2) * 0.125;
        let along_col =
            (5.0 * c + 4.0 * (n1 + s1) - diag1 - vert2 + 0.5 * horiz2) * 0.125;

        let red_horizontal = red_is_horizontal(gid.x, gid.y);
        let r = select(along_col, along_row, red_horizontal);
        let b = select(along_row, along_col, red_horizontal);
        rgb = vec3<f32>(r, c, b);
    } else {
        // ---- Red or blue site ----------------------------------------
        // Green at an R/B site: bilinear on the axial neighbours, corrected
        // by the centre channel's Laplacian.
        //   (4c + 2(n1+s1+w1+e1) - (n2+s2+w2+e2)) / 8
        let green = (4.0 * c + 2.0 * axial1 - (vert2 + horiz2)) * 0.125;

        // The opposite chroma sits on the diagonals.
        //   (6c + 2(nw+ne+sw+se) - 1.5(n2+s2+w2+e2)) / 8
        let opposite = (6.0 * c + 2.0 * diag1 - 1.5 * (vert2 + horiz2)) * 0.125;

        if (colour == 0u) {
            rgb = vec3<f32>(c, green, opposite);
        } else {
            rgb = vec3<f32>(opposite, green, c);
        }
    }

    // The correction term can overshoot below zero near clipped highlights.
    // Negative light is not meaningful, and carrying it forward makes the
    // ratio-based operations downstream (white balance, saturation) misbehave.
    rgb = max(rgb, vec3<f32>(0.0));

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