// G-buffer pack compute shader for CNN v3
// Pass 4: Pack all G-buffer channels into two rgba32uint feature textures (32 bytes/pixel)
// Output feat_tex0 holds 8×f16 geometric channels; feat_tex1 holds 12×u8 context channels.

struct GBufRes {
    resolution: vec2f,
}

@group(0) @binding(0) var<uniform> gbuf_res:      GBufRes;
@group(0) @binding(1) var gbuf_albedo:     texture_2d<f32>;
@group(0) @binding(2) var gbuf_normal_mat: texture_2d<f32>;
@group(0) @binding(3) var gbuf_depth:      texture_depth_2d;
@group(0) @binding(4) var gbuf_shadow:     texture_2d<f32>;
@group(0) @binding(5) var gbuf_transp:     texture_2d<f32>;
@group(0) @binding(6) var prev_cnn:        texture_2d<f32>;
@group(0) @binding(7) var feat_tex0:       texture_storage_2d<rgba32uint, write>;
@group(0) @binding(8) var feat_tex1:       texture_storage_2d<rgba32uint, write>;
@group(0) @binding(9) var bilinear_sampler: sampler;

// Sample depth texture at integer coordinate, clamp to borders.
fn load_depth(coord: vec2i) -> f32 {
    let dims = vec2i(textureDimensions(gbuf_depth));
    let c    = clamp(coord, vec2i(0), dims - vec2i(1));
    return textureLoad(gbuf_depth, c, 0);
}

// Box-filter albedo: average of 2×2 texels starting at top-left corner `tl`.
fn box2(tl: vec2i) -> vec3f {
    let a = textureLoad(gbuf_albedo, tl + vec2i(0, 0), 0).rgb;
    let b = textureLoad(gbuf_albedo, tl + vec2i(1, 0), 0).rgb;
    let c = textureLoad(gbuf_albedo, tl + vec2i(0, 1), 0).rgb;
    let d = textureLoad(gbuf_albedo, tl + vec2i(1, 1), 0).rgb;
    return (a + b + c + d) * 0.25;
}

// Box-filter albedo: average of 4×4 texels starting at top-left corner `tl`.
fn box4(tl: vec2i) -> vec3f {
    var acc = vec3f(0.0);
    for (var dy: i32 = 0; dy < 4; dy++) {
        for (var dx: i32 = 0; dx < 4; dx++) {
            acc += textureLoad(gbuf_albedo, tl + vec2i(dx, dy), 0).rgb;
        }
    }
    return acc * (1.0 / 16.0);
}

// Decode oct-normal from [0,1] storage → [-1,1] encoded xy → reconstruct z.
fn decode_oct_normal(rg: vec2f) -> vec3f {
    let f = rg * 2.0 - vec2f(1.0);
    var n = vec3f(f.x, f.y, 1.0 - abs(f.x) - abs(f.y));
    let t = max(-n.z, 0.0);
    n.x += select(t, -t, n.x >= 0.0);
    n.y += select(t, -t, n.y >= 0.0);
    return normalize(n);
}

@compute @workgroup_size(8, 8)
fn pack_features(@builtin(global_invocation_id) id: vec3u) {
    let coord = vec2i(id.xy);
    let dims  = vec2i(textureDimensions(gbuf_albedo));
    if (coord.x >= dims.x || coord.y >= dims.y) { return; }

    let uv = (vec2f(coord) + vec2f(0.5)) / gbuf_res.resolution;

    // --- Geometric channels (high precision, f16 packed) ---
    let albedo      = textureLoad(gbuf_albedo,     coord, 0).rgb;
    let nm          = textureLoad(gbuf_normal_mat, coord, 0);
    let depth_raw   = load_depth(coord);

    // Finite-difference depth gradient (central difference, clamped coords)
    let dzdx = (load_depth(coord + vec2i(1, 0)) - load_depth(coord - vec2i(1, 0))) * 0.5;
    let dzdy = (load_depth(coord + vec2i(0, 1)) - load_depth(coord - vec2i(0, 1))) * 0.5;

    // Normal: stored as oct-encoded [0,1] in RG; extract just the encoded xy for feat_tex0
    let normal_enc = nm.rg; // already in [0,1] — decode to get the xy for CNN input
    let n3         = decode_oct_normal(normal_enc);
    // Store oct-encoded in [-1,1] remapped back to what CNN expects (the [-1,1] oct xy)
    let oct_xy     = normal_enc * 2.0 - vec2f(1.0); // remap [0,1] → [-1,1]

    // Texture 0: 4 u32, each = pack2x16float of two f16 values
    // [0] albedo.r  | albedo.g
    // [1] albedo.b  | normal.x (oct, [-1,1])
    // [2] normal.y  | depth
    // [3] dzdx      | dzdy
    let t0 = vec4u(
        pack2x16float(albedo.rg),
        pack2x16float(vec2f(albedo.b, oct_xy.x)),
        pack2x16float(vec2f(oct_xy.y, depth_raw)),
        pack2x16float(vec2f(dzdx, dzdy))
    );
    textureStore(feat_tex0, coord, t0);

    // --- Context channels (low precision, u8 packed) ---
    let mat_id_u8 = nm.b;                                       // mat_id already in [0,1]
    let shadow    = textureLoad(gbuf_shadow, coord, 0).r;
    let transp    = textureLoad(gbuf_transp, coord, 0).r;
    let prev      = textureSampleLevel(prev_cnn, bilinear_sampler, uv, 0.0).rgb;

    // MIP 1: 2×2 box filter (half resolution context)
    // Use top-left aligned 2×2 block at half-res position
    let tl1  = coord * 2;  // this pixel's 2×2 region in full-res (mip1 is at half-res)
    // Actually we want to sample the neighborhood around this pixel for downsampled context.
    // mip1: sample a 2×2 box centered on the pixel in full-res coordinates
    let tl1c = max(coord - vec2i(0), vec2i(0));
    let mip1 = box2(tl1c);

    // mip2: sample a 4×4 box
    let tl2c = max(coord - vec2i(1), vec2i(0));
    let mip2 = box4(tl2c);

    // Texture 1: 4 u32, each = pack4x8unorm of four u8 values
    // [0] mat_id | prev.r | prev.g | prev.b
    // [1] mip1.r | mip1.g | mip1.b | mip2.r
    // [2] mip2.g | mip2.b | shadow | transp
    // [3] spare  (0)
    let t1 = vec4u(
        pack4x8unorm(vec4f(mat_id_u8, prev.r, prev.g, prev.b)),
        pack4x8unorm(vec4f(mip1.r, mip1.g, mip1.b, mip2.r)),
        pack4x8unorm(vec4f(mip2.g, mip2.b, shadow, transp)),
        0u
    );
    textureStore(feat_tex1, coord, t1);
}