// Hot and dead photosite repair, on the raw mosaic, before demosaic.
//
// A hot photosite reads far above anything the light put there — a leaky
// well, lit by its own dark current on a long or high-ISO exposure. Left in,
// the demosaic spreads it into its neighbours' interpolated channels and it
// becomes a coloured cross, three pixels wide, that no later stage can take
// back out: by then it is five pixels of plausible colour rather than one
// photosite of nonsense. So it is repaired here, where it is still one value.
//
// **What counts as hot.** A photosite far above *every* photosite of its own
// colour in its 5x5 window, and also far above every one of its eight
// immediate neighbours whatever their colour. The second half is what keeps a
// star or a glint: real light arrives through a lens and an anti-aliasing
// filter, so even the sharpest point lands on a patch of photosites, and the
// ones beside it are lit too. A hot photosite's neighbours are as dark as the
// rest of the frame. Dead photosites are the mirror image and are handled the
// same way.
//
// **What it becomes.** The brightest (for a hot photosite) or darkest (for a
// dead one) same-colour neighbour — the value nearest to what it read that
// the neighbourhood can vouch for. An average would soften the one case this
// gets wrong, a real highlight that happened to pass both tests; a clamp to
// the neighbourhood's range cannot invent anything.
//
// Written for either colour filter array: the colour of a photosite comes from
// a 6x6 tile anchored to the sensor, which holds the X-Trans pattern as it is
// and a Bayer 2x2 cell repeated nine times.

struct HotPixelParams {
    // The cropped area, in sensor photosites. Only photosites inside it are
    // judged, and only photosites inside it are asked as neighbours — the
    // masked border sits at black and would make everything look hot.
    crop_x: u32,
    crop_y: u32,
    width: u32,
    height: u32,
    // Row stride of the readout, in samples, and the number of u32 words.
    stride: u32,
    words: u32,
    // How many invocations one row of the dispatch grid holds, so a frame
    // wider than a dispatch dimension can be addressed as two.
    row_invocations: u32,
    // Samples in the readout. One less than twice `words` when the count is
    // odd, and the padding half of the last word is never judged.
    samples: u32,
    // Per-position black levels and reciprocal ranges, indexed by the
    // photosite's parity within the *crop*: (y&1)*2 + (x&1) counted from its
    // origin, as the demosaic counts them.
    black: vec4<f32>,
    inv_range: vec4<f32>,
    // The 6x6 colour tile, two bits per photosite, indexed by sensor
    // coordinate modulo 6: 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: HotPixelParams;
@group(0) @binding(2) var<storage, read_write> repaired: array<u32>;

// How far above its brightest neighbour a photosite must read to be hot, as a
// ratio and a margin in normalised units. Twice the neighbourhood and two
// percent of the range above it: far enough that shot noise in a lit area
// never qualifies, near enough that a hot photosite in a night sky — reading
// a third of the range over a sky at one percent — always does.
const HOT_RATIO: f32 = 2.0;
const HOT_MARGIN: f32 = 0.02;
// A dead photosite reads under half its darkest neighbour, and only counts
// where that neighbour is at least this bright: in the shadows, a photosite
// at zero is noise that clipped at the black point, not a defect.
const DEAD_RATIO: f32 = 0.5;
const DEAD_FLOOR: f32 = 0.05;

fn value_at(index: u32) -> u32 {
    let word = raw[index >> 1u];
    return select(word & 0xFFFFu, word >> 16u, (index & 1u) == 1u);
}

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;
}

// A raw value against its own black level and range. Compared rather than
// stored, so it is left unclamped at the top: a hot photosite above white is
// still more above white than its neighbours are.
fn level(sx: u32, sy: u32, v: u32) -> f32 {
    let cell = ((sy - params.crop_y) & 1u) * 2u + ((sx - params.crop_x) & 1u);
    return max(f32(v) - params.black[cell], 0.0) * params.inv_range[cell];
}

// The value to store for the photosite at `index`.
fn repair(index: u32) -> u32 {
    let v = value_at(index);
    let sx = index % params.stride;
    let sy = index / params.stride;
    if (sx < params.crop_x || sy < params.crop_y
        || sx >= params.crop_x + params.width || sy >= params.crop_y + params.height) {
        return v;
    }

    let centre = level(sx, sy, v);
    let colour = colour_at(sx, sy);

    var same_hi = -1.0;
    var same_lo = 1.0e9;
    var same_hi_raw = v;
    var same_lo_raw = v;
    var same_count = 0u;
    var adjacent_hi = 0.0;
    var adjacent_lo = 1.0e9;

    for (var dy = -2; dy <= 2; dy++) {
        for (var dx = -2; dx <= 2; dx++) {
            if (dx == 0 && dy == 0) {
                continue;
            }
            let nx = i32(sx) + dx;
            let ny = i32(sy) + dy;
            if (nx < i32(params.crop_x) || ny < i32(params.crop_y)
                || nx >= i32(params.crop_x + params.width)
                || ny >= i32(params.crop_y + params.height)) {
                continue;
            }
            let nsx = u32(nx);
            let nsy = u32(ny);
            let nv = value_at(nsy * params.stride + nsx);
            let n = level(nsx, nsy, nv);

            if (abs(dx) <= 1 && abs(dy) <= 1) {
                adjacent_hi = max(adjacent_hi, n);
                adjacent_lo = min(adjacent_lo, n);
            }
            if (colour_at(nsx, nsy) == colour) {
                same_count += 1u;
                if (n > same_hi) {
                    same_hi = n;
                    same_hi_raw = nv;
                }
                if (n < same_lo) {
                    same_lo = n;
                    same_lo_raw = nv;
                }
            }
        }
    }

    // A corner of the crop can leave a photosite with a single same-colour
    // neighbour, and one witness is not a neighbourhood.
    if (same_count < 2u) {
        return v;
    }

    let hot_line_same = same_hi * HOT_RATIO + HOT_MARGIN;
    let hot_line_adjacent = adjacent_hi * HOT_RATIO + HOT_MARGIN;
    if (centre > hot_line_same && centre > hot_line_adjacent) {
        return same_hi_raw;
    }
    if (same_lo >= DEAD_FLOOR && centre < same_lo * DEAD_RATIO
        && centre < adjacent_lo * DEAD_RATIO) {
        return same_lo_raw;
    }
    return v;
}

// One invocation per u32 word: two photosites, packed as the demosaic reads
// them. A word may straddle two rows when the stride is odd, which `repair`
// does not mind — it addresses by sample index.
@compute @workgroup_size(64, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let word = gid.y * params.row_invocations + gid.x;
    if (word >= params.words) {
        return;
    }
    let lo = repair(word * 2u);
    // The padding half of an odd-length readout is copied, not judged: it is
    // not a photosite, and the demosaic never addresses it.
    var hi = raw[word] >> 16u;
    if (word * 2u + 1u < params.samples) {
        hi = repair(word * 2u + 1u);
    }
    repaired[word] = (lo & 0xFFFFu) | (hi << 16u);
}