struct Params {
    block_size: vec2<i32>,
    offsets_stride: u32,
    search_radius: i32,
    skip_threshold: f32,
}

struct BlockOffset {
    score: f32,
    offset: vec2<i32>,
    tint: vec3<f32>,
}

@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var<storage, read_write> offsets: array<BlockOffset>;
@group(0) @binding(2) var next: texture_2d<f32>;
@group(0) @binding(3) var prev: texture_2d<f32>;

var<private> best_offset: vec2<i32> = vec2(0);
var<private> best_error: f32 = 100000.;
var<private> best_tint: vec3<f32> = vec3(0.);

@compute @workgroup_size(1)fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let uv = vec2<i32>(global_id.xy) * params.block_size;

    loop {
        test_offset(uv, vec2(0, 0));
        if best_error < params.skip_threshold { break; }
        apply_tint(uv);
        if best_error < params.skip_threshold { break; }
        best_tint = vec3(0.);

        do_dist(uv, 32);
        do_dist(uv, 16);
        do_dist(uv, 10);
        do_dist(uv, 8);
        if best_error < params.skip_threshold { break; }
        do_dist(uv, 6);
        do_dist(uv, 3);
        do_dist(uv, 4);
        do_dist(uv, 2);
        do_dist(uv, 1);
        if best_error < params.skip_threshold { break; }

        apply_tint(uv);
        break;
    }

    offsets[global_id.x + global_id.y * params.offsets_stride] = BlockOffset(best_error, best_offset, best_tint);
}

fn do_dist(uv: vec2<i32>, n: i32) {
    test_offset(uv, vec2(0, n));
    test_offset(uv, vec2(n, n));
    test_offset(uv, vec2(n, 0));
    test_offset(uv, vec2(n, -n));
    test_offset(uv, vec2(0, -n));
    test_offset(uv, vec2(-n, -n));
    test_offset(uv, vec2(-n, 0));
    test_offset(uv, vec2(-n, n));
}

fn apply_tint(uv: vec2<i32>) {

    var average_pcol = vec3(0.);
    var average_ncol = vec3(0.);
    for (var x = 0; x < params.block_size.x; x++) {
        for (var y = 0; y < params.block_size.y; y++) {
            let base = uv + vec2(x, y);
            average_pcol += textureLoad(prev, base + best_offset, 0).rgb;
            average_ncol += textureLoad(next, base, 0).rgb;
        }
    }

    let tint = (average_ncol - average_pcol) / f32(params.block_size.x * params.block_size.y);

    var err = 0.;
    for (var x = 0; x < params.block_size.x; x++) {
        for (var y = 0; y < params.block_size.y; y++) {
            let base = uv + vec2(x, y);
            let pcol = textureLoad(prev, base + best_offset, 0).rgb + tint;
            let ncol = textureLoad(next, base, 0).rgb;
            err += distance(pcol, ncol);
        }
    }
    if err < best_error {
        best_tint = tint;
    }
}

fn test_offset(uv: vec2<i32>, offset: vec2<i32>) {
    var err = 0.;
    for (var x = 0; x < params.block_size.x; x++) {
        for (var y = 0; y < params.block_size.y; y++) {
            let base = uv + vec2(x, y);
            let pcol = textureLoad(prev, base + offset, 0).rgb;
            let ncol = textureLoad(next, base, 0).rgb;
            err += distance(pcol, ncol);
        }
    }
    if err < best_error {
        best_error = err;
        best_offset = offset;
    }
}