#version 430
layout(local_size_x = 1, local_size_y = 1) in;

#include <utils>
#include <complex>

#define MAX_ITERATIONS 200

const float antialiasing = 0.28;
uniform float seconds = 0;
layout(rgba8) uniform writeonly image2D output_image;

vec2 get_step(vec2 z) {
  vec2 value = cpow(z, 5) + 2.0 * cpow(z, 3) - vec2(1., 0.);
  vec2 derivative = 5.0 * cpow(z, 4) + 3.0 * cpow(z, 2);
  return cdiv(value, derivative);
}

vec3 image(vec2 uv) {
  vec2 prevZ;
  vec2 z = uv * 3.0;

  float t = seconds / 20;
  vec2 factor = vec2(mod(t + 0.1, 2), 0);

  int iterations = 0;
  for (int i = 0; i < MAX_ITERATIONS; i++) {
    prevZ = z;
    z = prevZ - cmul(factor, get_step(prevZ));

    iterations++;
    if (abs(z.x - prevZ.x) < 0.001) break;
  }

  float val = float(iterations) * 10.0 / float(MAX_ITERATIONS);
  float r = 0.1;
  float g = min(val * 0.2, 0.7);
  float b = min(val * 0.8, 0.7);
  return srgb(r, g, b);
}

vec3 antialiased_image(vec2 uv) {
  vec2 resolution = vec2(imageSize(output_image).xy);
  float s = antialiasing / max(resolution.x, resolution.y);
  vec3 sum = vec3(0);
  int count = 0;
  for (int i = -2; i < 2; i++)
    for (int j = -2; j < 2; j++, count++)
      sum += image(uv + vec2(i, j) * s);
  return sum / count;
}

void main() {
  vec2 resolution = vec2(imageSize(output_image).xy);
  float aspect_ratio = resolution.y / resolution.x;
  vec2 uv = (gl_GlobalInvocationID.xy - resolution * 0.5) * vec2(1.0, aspect_ratio) / resolution;

  // vec3 color = image(uv);
  vec3 color = antialiased_image(uv);
  imageStore(output_image, ivec2(gl_GlobalInvocationID.xy), vec4(color, 1));
}