// 3x3 convolution with weight indexing
// Samples 9 pixels, applies mat4 weights per sample

fn cnn_conv3x3(
  tex: texture_2d<f32>,
  samp: sampler,
  uv: vec2<f32>,
  resolution: vec2<f32>,
  weights: array<mat4x4<f32>, 9>,
  bias: vec4<f32>
) -> vec4<f32> {
  let step = 1.0 / resolution;
  var sum = bias;
  var idx = 0;

  for (var dy = -1; dy <= 1; dy++) {
    for (var dx = -1; dx <= 1; dx++) {
      let offset = vec2<f32>(f32(dx), f32(dy)) * step;
      let sample = textureSample(tex, samp, uv + offset);
      sum += weights[idx] * sample;
      idx++;
    }
  }

  return sum;
}

fn cnn_conv3x3_with_coord(
  tex: texture_2d<f32>,
  samp: sampler,
  uv: vec2<f32>,
  resolution: vec2<f32>,
  rgba_weights: array<mat4x4<f32>, 9>,
  coord_weights: mat2x4<f32>,
  bias: vec4<f32>
) -> vec4<f32> {
  let step = 1.0 / resolution;
  var sum = bias;

  sum += coord_weights * uv;

  var idx = 0;
  for (var dy = -1; dy <= 1; dy++) {
    for (var dx = -1; dx <= 1; dx++) {
      let offset = vec2<f32>(f32(dx), f32(dy)) * step;
      let rgba = textureSample(tex, samp, uv + offset);
      sum += rgba_weights[idx] * rgba;
      idx++;
    }
  }

  return sum;
}