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

#include <sdf>
#include <utils>

uniform float seconds = 0;
layout(rgba8) uniform readonly image2D noise_texture;
layout(rgba8) uniform writeonly image2D output_image;
uniform vec3 camera_position = vec3(0, 1, 0);
uniform vec3 uv_direction = vec3(0);
uniform vec3 primary_light_position = vec3(1, 2, 1);

#define MAX_RAYMARCH_ITERATIONS 500
#define MAX_DISTANCE 100.0
#define SURFACE_DIST 0.01

// Noise Related Functions
float mod289(float x) { return x - floor(x * (1.0 / 289.0)) * 289.0; }
vec4 mod289(vec4 x) { return x - floor(x * (1.0 / 289.0)) * 289.0; }
vec4 perm(vec4 x) { return mod289(((x * 34.0) + 1.0) * x); }
float noise(vec3 p) {
  vec3 a = floor(p);
  vec3 d = p - a;
  d = d * d * (3.0 - 2.0 * d);
  vec4 b = a.xxyy + vec4(0.0, 1.0, 0.0, 1.0);
  vec4 k1 = perm(b.xyxy);
  vec4 k2 = perm(k1.xyxy + b.zzww);
  vec4 c = k2 + a.zzzz;
  vec4 k3 = perm(c);
  vec4 k4 = perm(c + 1.0);
  vec4 o1 = fract(k3 * (1.0 / 41.0));
  vec4 o2 = fract(k4 * (1.0 / 41.0));
  vec4 o3 = o2 * d.z + o1 * (1.0 - d.z);
  vec2 o4 = o3.yw * d.x + o3.xz * (1.0 - d.x);
  return o4.y * d.y + o4.x * (1.0 - d.y);
}
float PHI = 1.61803398874989484820459; // Φ = Golden Ratio
float gold_noise(in vec2 xy, in float seed) {
  return fract(tan(distance(xy * PHI, xy) * seed) * xy.x);
}
float random(vec3 scale, float seed) {
  return fract(sin(dot(gl_GlobalInvocationID.xyz + seed, scale)) * 43758.5453 +
               seed);
}

// Fractal Brownian Noise
float fbm(vec3 x, int octaves) {
  float v = 0.0;
  float a = 0.5;
  vec3 shift = vec3(100);
  for (int i = 0; i < octaves; ++i) {
    v += a * noise(x);
    x = x * 2.0 + shift;
    a *= 0.5;
  }
  return v;
}

float random(vec2 st) {
  return fract(sin(dot(st.xy, vec2(12.9898, 78.233))) * 43758.5453123);
}

float noise(vec2 st) {
  vec2 i = floor(st);
  vec2 f = fract(st);
  float a = random(i);
  float b = random(i + vec2(1.0, 0.0));
  float c = random(i + vec2(0.0, 1.0));
  float d = random(i + vec2(1.0, 1.0));
  vec2 u = f * f * (3.0 - 2.0 * f);
  return mix(a, b, u.x) + (c - a) * u.y * (1.0 - u.x) + (d - b) * u.x * u.y;
}

float fbm(vec2 st) {
  float value = 0.0;
  float amplitude = 0.5;
  for (int i = 0; i < 5; i++) {
    value += amplitude * noise(st);
    st *= 2.0;
    amplitude *= 0.5;
  }
  return value;
}

// TODO: Enumify
#define OID_TERRAIN 4.

vec2 scene(in vec3 p) {
  float c1 = sdf_circle(p - vec3(0, 1, 2.4), 0.2);
  float c2 = sdf_circle(p - vec3(-0.6, 1, 3), 0.3);
  float box = sdf_box(p - vec3(1, 0.5, 3), vec3(0.5, 0.18, 0.2), 0.07);
  float terrain_noise = 6.0 * noise(p / 20.0) + 0.2 * noise((p + 200.0) / 2.0) +
                        0.005 * noise((p - 50.0) * 20.0);
  float terrain =
      sdf_plane(p + vec3(0.0, terrain_noise, 0.0), vec3(0., 1., 0.), 0.0);
  float light_indicator =
      sdf_circle(p - primary_light_position + vec3(0.1, 0.1, -0.4), 0.01);
  float dist = min(min(min(min(c1, c2), terrain), box), light_indicator);
  float oid = 0.;
  if (dist == c1)
    oid = 1.;
  else if (dist == c2)
    oid = 2.;
  else if (dist == box)
    oid = 3.;
  else if (dist == terrain)
    oid = OID_TERRAIN;
  else if (dist == light_indicator)
    oid = -1.;
  return vec2(dist, oid);
}

vec2 ray_march_scene(in vec3 ray_origin, in vec3 ray_direction) {
  float dist = 0.;
  for (int i = 0; i < MAX_RAYMARCH_ITERATIONS && dist < MAX_DISTANCE; i++) {
    vec3 p = ray_origin + dist * ray_direction;
    vec2 obj = scene(p);
    float gap = obj.x;
    dist += gap;
    if (gap <= SURFACE_DIST)
      return vec2(dist, obj.y);
  }
  return vec2(dist, 0.0);
}

vec3 surface_normal_scene(in vec3 p) {
  vec2 inc = vec2(0.01, 0);
  return normalize(vec3(scene(p + inc.xyy).x - scene(p - inc.xyy).x,
                        scene(p + inc.yxy).x - scene(p - inc.yxy).x,
                        scene(p + inc.yyx).x - scene(p - inc.yyx).x));
}

float shadow_scene(in vec3 ray_origin, in vec3 ray_direction,
                   in float light_intensity) {
  float dist = 0.;
  float res = 1.;
  float prev_gap = 1e10;
  for (int i = 0; i < 50 && dist < MAX_DISTANCE; i++) {
    vec3 p = ray_origin + dist * ray_direction;
    vec2 obj = scene(p);
    float gap = obj.x;
    if (obj.y < 0.) // Debug objects dont cast shadows
      continue;
    float y = gap * gap / (10000.0 * prev_gap);
    float d = sqrt(abs(gap * gap - y * y));
    res = min(res, light_intensity * d / max(0., dist - y));
    prev_gap = gap;
    dist += gap;

    if (res < 0.0001)
      break;
  }
  res = clamp(res, 0.08, 1.0);

  return res * res * (3.0 - 2.0 * res);
}

void main() {
  vec2 resolution = vec2(imageSize(output_image).xy);
  vec3 uv = vec3((gl_GlobalInvocationID.xy - resolution / 2.0) / resolution, 0);
  float aspect_ratio = resolution.y / resolution.x;
  uv /= vec3(aspect_ratio, 1, 1);

  vec3 uv_direction = normalize(vec3(uv.x, uv.y, 1));

  vec3 light_pos = primary_light_position;

  vec2 obj = ray_march_scene(camera_position, uv_direction);
  float depth = obj.x;
  float oid = obj.y;
  vec3 p = camera_position + uv_direction * depth;
  vec3 normal = surface_normal_scene(p);
  vec3 light_dir = normalize(light_pos - p);
  vec3 hal = normalize(light_dir - uv_direction);

  vec3 col = vec3(0.9, 0.3, 0.4);
  if (oid < 0.) // Debug objects
    col = vec3(1.0, 0.2, 0.2);
  else if (oid == OID_TERRAIN)
    col = vec3(0.24, 0.61, 0.08);
  else if (oid == 0) { // Sky
    vec4 cloud_noise =
        imageLoad(noise_texture, ivec2(gl_GlobalInvocationID.xy));
    col = vec3(0.35, 0.6, 0.9);
    col = mix(col, cloud_noise.xyz, clamp(0.5, 0., 1.));
  }

  if (depth < MAX_DISTANCE && oid > 0.) {
    float contrast = 1.;
    contrast *= clamp(dot(normal, light_dir), 0, 1.);
    contrast *= shadow_scene(p, light_dir, 24.0);
    float spe =
        pow(clamp(dot(normal, hal), 0.0, 1.0), 16.0) * contrast *
        (0.04 + 0.96 * pow(clamp(1.0 + dot(hal, uv_direction), 0.0, 1.0), 5.0));
    col = col * 1.1 + 0.01;
    if (oid != OID_TERRAIN)
      col += 12.0 * spe * vec3(1.00, 0.70, 0.5);
    col *= contrast;
    // col *= max(0.1, exp(-0.00005 * depth * depth * depth));
  }

  imageStore(output_image, ivec2(gl_GlobalInvocationID.xy),
             vec4(clamp(col, 0., 1.), 1));
}