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

#include <sdf>
#include <utils>
#include <noise>

#define RM_MAX_ITERATIONS 400
#define RM_SHADOW_MAX_ITERATIONS 60
#define RM_MAX_DISTANCE 100.
#define DEFAULT_BLEND_GAP 0.001

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 mat3 camera_transformation;
uniform vec3 key_light_position = vec3(1, 2, 1);
uniform vec3 fill1_light_position = vec3(-1, 1.7, -1);
uniform vec3 fill2_light_position = vec3(1, 1.6, 1);

struct Obj {
  int object_id;
  float dist;
  vec3 material;
  bool lighting;
};

Obj sky_object = Obj(-1, 0., vec3(0.35, 0.6, 0.9), false);

Obj create_object(in int object_id, in float dist) {
  return Obj(object_id, dist, vec3(0), true);
}

#define OBJ_FACE 1
#define OBJ_EYES 2
#define OBJ_HAIR 4

float sample_noise(vec2 p) {
  vec4 v = imageLoad(noise_texture, ivec2(p));
  return v.x;
}

Obj scene(in vec3 p) {
  float head = sdf_circle(p - vec3(-0.1, 1.1, 1.0), 0.42);
  float jaw = sd_ellipsoid(p - vec3(-0.1, 0.86, 0.86), vec3(0.36, 0.36, 0.2));
  float eyel_socket = sdf_circle(p - vec3(-0.3, 1.1, 0.57), 0.118);
  float eyer_socket = sdf_circle(p - vec3(0.05, 1.1, 0.56), 0.118);

  float mouth_rot = -0.5;
  vec3 mouthp = p - vec3(-0.1, 0.96, 0.68);
  mouthp.yz *= mat2x2(cos(mouth_rot), -sin(mouth_rot), sin(mouth_rot), cos(mouth_rot));
  float mouth = sd_ellipsoid(mouthp, vec3(0.1, 0.08, 0.04));
  Obj face = create_object(OBJ_FACE,
      ssubtract(
        ssubtract(
          smin(head, jaw, 0.12),
          min(eyel_socket, eyer_socket), 0.06),
        mouth, 0.07));
  face.material = vec3(0.99, 0.64, 0.61);

  float eyel = sdf_circle(p - vec3(-0.26, 1.1, 0.68), 0.08);
  float eyer = sdf_circle(p - vec3(0.03, 1.1, 0.67), 0.08);
  float dirx = length(normalize(p.xz + vec2(0.5, 0.0)));
  Obj eyes = create_object(OBJ_EYES, min(eyel, eyer));

  Obj teeth = create_object(3,
      min(sdf_box(p - vec3(-0.08, 0.95, 0.61), vec3(0.01, 0.02, 0.001), 0),
        sdf_box(p - vec3(-0.1, 0.955, 0.61), vec3(0.01, 0.02, 0.001), 0)));
  teeth.material = vec3(1);

  Obj hair = create_object(OBJ_HAIR,
      ssubtract(sdf_circle(p - vec3(-0.1, 1.16, 1.036) * (1.0 + (fbm(p.xy * 10) - 0.1) / 5 + sin(100 * p.x) / 200), 0.44),
        sdf_circle(p - vec3(-0.1, 1.1, 1.), 0.4), 0.08));
  hair.material = vec3(0.02, 0.01, 0.01);

  Obj platform = create_object(999, sdf_box(p - vec3(0, 0.1, 5), vec3(10, 0.01, 10), 0.0));
  platform.material = vec3(1);

  float dist = platform.dist;
  dist = min(dist, face.dist);
  dist = min(dist, eyes.dist);
  dist = min(dist, teeth.dist);
  dist = min(dist, hair.dist);

  if (dist == eyes.dist) return eyes;
  if (dist == teeth.dist) return teeth;
  if (dist == face.dist) return face;
  if (dist == hair.dist) return hair;
  if (dist == platform.dist) return platform;
  return sky_object; // unreachable
}

Obj ray_march_scene(in vec3 ray_origin, in vec3 ray_direction) {
  float dist = 0.;
  for (int i = 0; i < RM_MAX_ITERATIONS && dist < RM_MAX_DISTANCE; i++) {
    vec3 p = ray_origin + ray_direction * dist;
    Obj obj = scene(p);
    dist += obj.dist;
    if (obj.dist <= 0.001) {
      obj.dist = dist;
      return obj;
    }
  }
  return sky_object;
}

float shadow_scene(in vec3 ray_origin, in vec3 ray_direction, in float light_intensity) {
  float dist = 0.015;
  float res = 1.;
  float prev_gap = 1e10;
  for (int i = 0; i < RM_SHADOW_MAX_ITERATIONS && dist < RM_MAX_DISTANCE; i++) {
    vec3 p = ray_origin + dist * ray_direction;
    Obj obj = scene(p);
    float gap = obj.dist;
    if (!obj.lighting)
      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);
}

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

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 dir = vec3(uv.x, uv.y, 1);
  vec3 origin = vec3(0);
  vec3 uv_direction = normalize((dir - origin) * camera_transformation + origin);
  vec3 light_position = key_light_position;

  Obj obj = ray_march_scene(camera_position, uv_direction);
  vec3 p = camera_position + uv_direction * obj.dist;
  vec3 normal = surface_normal_scene(p);

  vec3 color = obj.material;
  if (obj.object_id == OBJ_EYES) {
    float foo = -dot(normal, vec3(0, 0, 1));
    vec3 iris = vec3(1);
    if (foo > 0.7) iris = mix(vec3(0.5, 0.3, 0.3), vec3(0), foo * 1.05);
    color = iris;
  }

  if (obj.lighting) {
    vec3 light_direction = normalize(light_position - p);
    float light_distance = length(light_position - p);
    float light_intensity = 32.0 / max(light_distance, 0.01);

    float contrast = 1.0;
    contrast *= clamp(dot(normal, light_direction), 0., 1.);
    contrast *= shadow_scene(p, light_direction, light_intensity);
    contrast = clamp(contrast, 0., 1.);

    color *= contrast;
    if (obj.object_id != OBJ_EYES && obj.object_id != OBJ_HAIR) {
      color += vec3(1.0, 0.5, 0.5) * shadow_scene(p, normalize(fill1_light_position - p), 0.4);
      color += vec3(1.0, 0.6, 0.6) * shadow_scene(p, normalize(fill2_light_position - p), 0.4);
    }
  }

  // color = vec3(pow(color.x, 2.2), pow(color.y, 2.2), pow(color.z, 2.2));
  color = clamp(color, 0., 1.);
  imageStore(output_image, ivec2(gl_GlobalInvocationID.xy), vec4(color, 1));
}