// CNN v3 — Encoder level 1
// AvgPool2x2(enc0) + Conv(8->16, 3x3, zero-pad) + FiLM + ReLU
//
// Input:  enc0_tex     (rgba32uint, 8xf16)  full-res
// Output: enc1_out_lo  (rgba32uint, 8xf16)  half-res  ch 0-7
//         enc1_out_hi  (rgba32uint, 8xf16)  half-res  ch 8-15
//
// Weight layout (f16, OIHW + bias):
//   [0 .. 8*16*9)   conv: w[out][in][ky][kx]
//   [1152 .. +16)   bias: b[out]

#include "cnn_v3/common"

const ENC1_IN:  u32 = 8u;
const ENC1_OUT: u32 = 16u;

struct Params {
    weight_offset: u32,
    _pad: vec3u,
    gamma_0: vec4f,
    gamma_1: vec4f,
    gamma_2: vec4f,
    gamma_3: vec4f,
    beta_0:  vec4f,
    beta_1:  vec4f,
    beta_2:  vec4f,
    beta_3:  vec4f,
}

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

fn film_gamma(o: u32) -> f32 {
    if (o < 4u)  { return params.gamma_0[o]; }
    if (o < 8u)  { return params.gamma_1[o - 4u]; }
    if (o < 12u) { return params.gamma_2[o - 8u]; }
    return params.gamma_3[o - 12u];
}
fn film_beta(o: u32) -> f32 {
    if (o < 4u)  { return params.beta_0[o]; }
    if (o < 8u)  { return params.beta_1[o - 4u]; }
    if (o < 12u) { return params.beta_2[o - 8u]; }
    return params.beta_3[o - 12u];
}

fn load_enc0_avg(hcoord: vec2i, full_dims: vec2i) -> array<f32, 8> {
    let half_dims = full_dims / 2;
    if (hcoord.x < 0 || hcoord.y < 0 || hcoord.x >= half_dims.x || hcoord.y >= half_dims.y) {
        return array<f32, 8>(0., 0., 0., 0., 0., 0., 0., 0.);
    }
    let base = hcoord * 2;
    var s: array<f32, 8>;
    for (var dy: i32 = 0; dy < 2; dy++) {
        for (var dx: i32 = 0; dx < 2; dx++) {
            let fc = clamp(base + vec2i(dx, dy), vec2i(0), full_dims - vec2i(1));
            let f  = unpack_8ch(enc0_tex, fc);
            for (var i: u32 = 0u; i < 8u; i++) { s[i] += f[i]; }
        }
    }
    for (var i: u32 = 0u; i < 8u; i++) { s[i] *= 0.25; }
    return s;
}

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

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

    for (var o: u32 = 0u; o < ENC1_OUT; o++) {
        var sum = get_w(wo, ENC1_OUT * ENC1_IN * 9u + o);  // bias
        for (var ky: i32 = -1; ky <= 1; ky++) {
            for (var kx: i32 = -1; kx <= 1; kx++) {
                let feat = load_enc0_avg(coord + vec2i(kx, ky), full_dims);
                let ki   = u32(ky + 1) * 3u + u32(kx + 1);
                for (var i: u32 = 0u; i < ENC1_IN; i++) {
                    sum += get_w(wo, o * ENC1_IN * 9u + i * 9u + ki) * feat[i];
                }
            }
        }
        out[o] = max(0.0, film_gamma(o) * sum + film_beta(o));
    }

    textureStore(enc1_out_lo, coord, vec4u(
        pack2x16float(vec2f(out[0],  out[1])),
        pack2x16float(vec2f(out[2],  out[3])),
        pack2x16float(vec2f(out[4],  out[5])),
        pack2x16float(vec2f(out[6],  out[7]))
    ));
    textureStore(enc1_out_hi, coord, vec4u(
        pack2x16float(vec2f(out[8],  out[9])),
        pack2x16float(vec2f(out[10], out[11])),
        pack2x16float(vec2f(out[12], out[13])),
        pack2x16float(vec2f(out[14], out[15]))
    ));
}