struct Settings {
  inverseDiffusionRateTrails: f32,
  decayRateTrails: f32,
  inverseDiffusionRateBrush: f32,
  decayRateBrush: f32,
  diffusionNeighborDivisor: f32,
  brushDecayAlphaOffset: f32,
  padding0: f32,
  padding1: f32,
};

const WORKGROUP_SIZE_X = 16u;
const WORKGROUP_SIZE_Y = 16u;
const TILE_SIZE_X = WORKGROUP_SIZE_X + 2u;
const TILE_SIZE_Y = WORKGROUP_SIZE_Y + 2u;
const TILE_TEXEL_COUNT = TILE_SIZE_X * TILE_SIZE_Y;

@group(0) @binding(0) var<uniform> settings: Settings;
@group(0) @binding(1) var trailMap: texture_2d<f32>;
@group(0) @binding(2) var trailMapOut: texture_storage_2d<rgba16float, write>;

var<workgroup> tile: array<vec4<f32>, 324>;

@compute @workgroup_size(16, 16)
fn main(
  @builtin(global_invocation_id) global_id: vec3<u32>,
  @builtin(local_invocation_id) local_id: vec3<u32>,
  @builtin(workgroup_id) workgroup_id: vec3<u32>
) {
  let textureSize = vec2<i32>(textureDimensions(trailMap, 0));
  let localLinearIndex = local_id.y * WORKGROUP_SIZE_X + local_id.x;
  var tileIndex = localLinearIndex;

  loop {
    if tileIndex >= TILE_TEXEL_COUNT {
      break;
    }

    let tilePosition = vec2<u32>(tileIndex % TILE_SIZE_X, tileIndex / TILE_SIZE_X);
    let sourcePixelU32 =
      workgroup_id.xy * vec2<u32>(WORKGROUP_SIZE_X, WORKGROUP_SIZE_Y) +
      tilePosition;
    let sourcePixel = clamp(
      vec2<i32>(i32(sourcePixelU32.x), i32(sourcePixelU32.y)) - vec2<i32>(1, 1),
      vec2<i32>(0, 0),
      textureSize - vec2<i32>(1, 1)
    );
    tile[tileIndex] = textureLoad(trailMap, sourcePixel, 0);
    tileIndex += WORKGROUP_SIZE_X * WORKGROUP_SIZE_Y;
  }

  workgroupBarrier();

  let pixel = vec2<i32>(i32(global_id.x), i32(global_id.y));
  if pixel.x >= textureSize.x || pixel.y >= textureSize.y {
    return;
  }

  let centerTilePosition = local_id.xy + vec2<u32>(1u, 1u);
  let centerTileIndex = centerTilePosition.y * TILE_SIZE_X + centerTilePosition.x;
  var current = tile[centerTileIndex];
  let random = random_from_pixel(pixel);
  let trailWeight = diffusion_weight(random, settings.inverseDiffusionRateTrails);
  let brushWeight = diffusion_weight(random, settings.inverseDiffusionRateBrush);

  current += (
      propagate(centerTileIndex, -1, -1, current, trailWeight, brushWeight)
    + propagate(centerTileIndex, -1, 1, current, trailWeight, brushWeight)
    + propagate(centerTileIndex, 1, -1, current, trailWeight, brushWeight)
    + propagate(centerTileIndex, 1, 1, current, trailWeight, brushWeight)

    + propagate(centerTileIndex, -1, 0, current, trailWeight, brushWeight)
    + propagate(centerTileIndex, 0, -1, current, trailWeight, brushWeight)
    + propagate(centerTileIndex, 1, 0, current, trailWeight, brushWeight)
    + propagate(centerTileIndex, 0, 1, current, trailWeight, brushWeight)
  ) / max(1.0, settings.diffusionNeighborDivisor);

  let decayed = clamp(vec4(
    current.rgb * settings.decayRateTrails,
    max(0, current.a + (current.a - settings.brushDecayAlphaOffset) * settings.decayRateBrush)
  ), vec4(0), vec4(1));

  textureStore(trailMapOut, pixel, decayed);
}

fn propagate(
  centerTileIndex: u32,
  offsetX: i32,
  offsetY: i32,
  currentColor: vec4<f32>,
  trailWeight: f32,
  brushWeight: f32
) -> vec4<f32> {
  let neighbourIndex = i32(centerTileIndex) + offsetY * i32(TILE_SIZE_X) + offsetX;
  let neighbour = tile[u32(neighbourIndex)];
  let difference = clamp(neighbour - currentColor, vec4(0), vec4(1));

  return vec4(
    vec3(length(neighbour.rgb) * trailWeight),
    neighbour.a * brushWeight
  ) * difference;
}

fn random_from_pixel(pixel: vec2<i32>) -> f32 {
  let p = vec2<u32>(pixel);
  var hash = p.x * 1664525u + p.y * 1013904223u + 374761393u;
  hash = (hash ^ (hash >> 16u)) * 2246822519u;
  hash = (hash ^ (hash >> 13u)) * 3266489917u;
  hash = hash ^ (hash >> 16u);
  return f32(hash) * 2.3283064365386963e-10;
}

fn diffusion_weight(random: f32, inverseRate: f32) -> f32 {
  let r = clamp(random, 0.0, 1.0);

  if inverseRate < 1.0 {
    let rootApproximation = r / max(0.5 + r * 0.5, 0.0001);
    return mix(
      rootApproximation,
      r,
      clamp((inverseRate - 0.5) * 2.0, 0.0, 1.0)
    );
  }

  let r2 = r * r;
  if inverseRate < 2.0 {
    return mix(r, r2, inverseRate - 1.0);
  }

  let r4 = r2 * r2;
  if inverseRate < 4.0 {
    return mix(r2, r4, (inverseRate - 2.0) * 0.5);
  }

  let r8 = r4 * r4;
  if inverseRate < 8.0 {
    return mix(r4, r8, (inverseRate - 4.0) * 0.25);
  }

  let r16 = r8 * r8;
  return mix(r8, r16, clamp((inverseRate - 8.0) * 0.125, 0.0, 1.0))
    * min(1.0, 16.0 / inverseRate);
}