// Copyright 2026 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// This material implements a single pass of the Jump Flooding Algorithm (JFA).
// This algorithm is used to efficiently compute a distance transform where each
// pixel stores the UV coordinates of its closest edge point. In each pass,
// pixels look at their neighbors at a distance 'step' to see if they have a
// closer edge point.
material {
    name : outline_jumpflood,
    culling : none,
    depthWrite : false,
    depthCulling : false,
    parameters : [
        { type : sampler2d, name : source },
        { type : float, name : step }
    ],
    variables : [
        vertex
    ],
    domain : postprocess
}

vertex {
  void postProcessVertex(inout PostProcessVertexInputs postProcess) {
    postProcess.vertex.xy = uvToRenderTargetUV(postProcess.normalizedUV);
  }
}

fragment {
  // Encode a vec2 value with the range [0-1] into a 32-bit (RGBA) value.
  // We do this by packing each coordinate into a 16-bit value. Values are
  // adjusted to ensure [0x0000] is reserved for "untouched" pixels rather
  // than the pixel at the position (0, 0).
  vec4 pack(vec2 value) {
    // Scale and bias the value to avoid (0,0,0,0).
    vec2 adjusted = value * 0.999 + 0.001;

    // Pack float [0, 1] into two 8-bit channels.
    float val_x = adjusted.x * 255.0;
    float r_x = floor(val_x) / 255.0;
    float g_x = fract(val_x);

    float val_y = adjusted.y * 255.0;
    float r_y = floor(val_y) / 255.0;
    float g_y = fract(val_y);

    return vec4(r_x, g_x, r_y, g_y);
  }

  vec2 unpack(vec4 value) {
    float x = value.r + value.g / 255.0;
    float y = value.b + value.a / 255.0;
    // Reverse scale and bias
    return (vec2(x, y) - 0.001) / 0.999;
  }

  vec3 jump(vec3 best, vec2 offset) {
    vec2 uv = variable_vertex.xy + offset;
    if (uv.x < 0.0 || uv.x > 1.0 || uv.y < 0.0 || uv.y > 1.0) {
      return best;
    }

    vec4 tex = texture(materialParams_source, uv);
    if (tex == vec4(0.0)) {
      return best;
    }

    vec2 edge = unpack(tex);
    vec2 diff = (edge - variable_vertex.xy);
    float dist_sq = dot(diff, diff);
    if (dist_sq < best.z) {
      best = vec3(edge.x, edge.y, dist_sq);
    }
    return best;
  }

  void postProcess(inout PostProcessInputs postProcess) {
    vec2 size = vec2(textureSize(materialParams_source, 0));
    vec2 dist = materialParams.step / size;

    vec3 best = vec3(-1.0, -1.0, 1e20);
    best = jump(best, dist * vec2(+1, +1));
    best = jump(best, dist * vec2(+1,  0));
    best = jump(best, dist * vec2(+1, -1));
    best = jump(best, dist * vec2( 0, +1));
    best = jump(best, dist * vec2( 0,  0));
    best = jump(best, dist * vec2( 0, -1));
    best = jump(best, dist * vec2(-1, +1));
    best = jump(best, dist * vec2(-1,  0));
    best = jump(best, dist * vec2(-1, -1));

    if (best.x >= 0.0) {
      postProcess.color = pack(best.xy);
    } else {
      postProcess.color = vec4(0.0);
    }
  }
}