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;

    test_offset(uv, vec2(0, 0));

    do_dist(uv, 32);
    do_dist(uv, 16);
    do_dist(uv, 10);
    do_dist(uv, 8);
    do_dist(uv, 6);
    do_dist(uv, 3);
    do_dist(uv, 4);
    do_dist(uv, 2);
    do_dist(uv, 1);

    apply_tint(uv);

    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 search_offset(uv: vec2<i32>) -> SearchRes {
//     var best_err = 100000000.;
//     var best_offset = vec2(0);
//     // TODO: better ordering
//     for (var ox = -params.search_radius; ox <= params.search_radius; ox++) {
//         for (var oy = -params.search_radius; oy <= params.search_radius; oy++) {
//             let offset = vec2(ox, oy);
//         }
//     }
//     return SearchRes(best_offset, best_err);
// }


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


// fn colormap_vec(v: vec2<f32>) -> vec3<f32> {
//     return vec3(v.y, v.x - 0.5 * v.y, -v.x - 0.5 * v.y);
// }