// CNN v3 — Decoder level 0 + output
// NearestUp2x(dec1) + cat(enc0_skip) -> Conv(8->4, 3x3, zero-pad) + FiLM + ReLU + Sigmoid
//
// Inputs: dec1_tex  (rgba16float, 4ch)  half-res
//         enc0_tex  (rgba16float, 4ch)  full-res  (skip connection)
// Output: output_tex (rgba16float, 4ch) full-res  (dispatch at full-res dims)
//
// Weight layout (f16, OIHW + bias):
//   [0 .. 8*4*9)  conv: w[out][in][ky][kx]   (in=8: 4 dec1 + 4 enc0 skip)
//   [288 .. +4)   bias: b[out]
//
// Parity note: sigmoid applied directly to dec0 output (matches train_cnn_v3.py forward()).

#include "cnn_v3/common"

const DEC0_IN:  u32 = 8u;
const DEC0_OUT: u32 = 4u;

struct Params {
    weight_offset: u32,
    _pad: vec3u,
    gamma: vec4f,
    beta:  vec4f,
}

@group(0) @binding(0) var dec1_tex: texture_2d<f32>;
@group(0) @binding(1) var enc0_tex: texture_2d<f32>;
@group(0) @binding(2) var<storage, read> weights: array<u32>;
@group(0) @binding(3) var<uniform> params: Params;
@group(0) @binding(4) var output_tex: texture_storage_2d<rgba16float, write>;

// Load 8 concatenated channels at full-res coord:
//   ch 0-3:  dec1 nearest-up (dec1_tex[coord/2])
//   ch 4-7:  enc0 skip       (enc0_tex[coord])
// Returns zeros for OOB coord (zero-padding for the conv).
fn load_dec0_concat(coord: vec2i, full_dims: vec2i) -> array<f32, 8> {
    if (coord.x < 0 || coord.y < 0 || coord.x >= full_dims.x || coord.y >= full_dims.y) {
        return array<f32, 8>(0., 0., 0., 0., 0., 0., 0., 0.);
    }
    let half_dims = vec2i(textureDimensions(dec1_tex));
    let hc  = clamp(coord / 2, vec2i(0), half_dims - vec2i(1));
    let d   = textureLoad(dec1_tex, hc, 0);
    let e   = textureLoad(enc0_tex, coord, 0);
    return array<f32, 8>(d.x, d.y, d.z, d.w, e.x, e.y, e.z, e.w);
}

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

    let wo = params.weight_offset;
    var out: array<f32, DEC0_OUT>;

    for (var o: u32 = 0u; o < DEC0_OUT; o++) {
        var sum = get_w(wo, DEC0_OUT * DEC0_IN * 9u + o);  // bias
        for (var ky: i32 = -1; ky <= 1; ky++) {
            for (var kx: i32 = -1; kx <= 1; kx++) {
                let feat = load_dec0_concat(coord + vec2i(kx, ky), full_dims);
                let ki   = u32(ky + 1) * 3u + u32(kx + 1);
                for (var i: u32 = 0u; i < DEC0_IN; i++) {
                    sum += get_w(wo, o * DEC0_IN * 9u + i * 9u + ki) * feat[i];
                }
            }
        }
        // FiLM + ReLU + Sigmoid (matches training forward())
        let v = max(0.0, params.gamma[o] * sum + params.beta[o]);
        out[o] = 1.0 / (1.0 + exp(-v));
    }

    textureStore(output_tex, coord, vec4f(out[0], out[1], out[2], out[3]));
}