// ray_tracer.cm — Simple ray tracer kernel
// Control flow + numeric + allocation. Classic VM benchmark.

var math = use('math/radians')

function vec(x, y, z) {
  return {x: x, y: y, z: z}
}

function vadd(a, b) {
  return {x: a.x + b.x, y: a.y + b.y, z: a.z + b.z}
}

function vsub(a, b) {
  return {x: a.x - b.x, y: a.y - b.y, z: a.z - b.z}
}

function vmul(v, s) {
  return {x: v.x * s, y: v.y * s, z: v.z * s}
}

function vdot(a, b) {
  return a.x * b.x + a.y * b.y + a.z * b.z
}

function vnorm(v) {
  var len = math.sqrt(vdot(v, v))
  if (len == 0) return vec(0, 0, 0)
  return vmul(v, 1 / len)
}

function make_sphere(center, radius, color) {
  return {
    center: center,
    radius: radius,
    color: color
  }
}

function intersect_sphere(origin, dir, sphere) {
  var oc = vsub(origin, sphere.center)
  var b = vdot(oc, dir)
  var c = vdot(oc, oc) - sphere.radius * sphere.radius
  var disc = b * b - c
  if (disc < 0) return -1
  var sq = math.sqrt(disc)
  var t1 = -b - sq
  var t2 = -b + sq
  if (t1 > 0.001) return t1
  if (t2 > 0.001) return t2
  return -1
}

function make_scene() {
  var spheres = [
    make_sphere(vec(0, -1, 5), 1, vec(1, 0, 0)),
    make_sphere(vec(2, 0, 6), 1, vec(0, 1, 0)),
    make_sphere(vec(-2, 0, 4), 1, vec(0, 0, 1)),
    make_sphere(vec(0, 1, 4.5), 0.5, vec(1, 1, 0)),
    make_sphere(vec(1, -0.5, 3), 0.3, vec(1, 0, 1)),
    make_sphere(vec(0, -101, 5), 100, vec(0.5, 0.5, 0.5))
  ]
  var light = vnorm(vec(1, 1, -1))
  return {spheres: spheres, light: light}
}

function trace(origin, dir, scene) {
  var closest_t = 999999
  var closest_sphere = null
  var i = 0
  var t = 0
  for (i = 0; i < length(scene.spheres); i++) {
    t = intersect_sphere(origin, dir, scene.spheres[i])
    if (t > 0 && t < closest_t) {
      closest_t = t
      closest_sphere = scene.spheres[i]
    }
  }

  if (!closest_sphere) return vec(0.2, 0.3, 0.5) // sky color

  var hit = vadd(origin, vmul(dir, closest_t))
  var normal = vnorm(vsub(hit, closest_sphere.center))
  var diffuse = vdot(normal, scene.light)
  if (diffuse < 0) diffuse = 0

  // Shadow check
  var shadow_origin = vadd(hit, vmul(normal, 0.001))
  var in_shadow = false
  for (i = 0; i < length(scene.spheres); i++) {
    if (scene.spheres[i] != closest_sphere) {
      t = intersect_sphere(shadow_origin, scene.light, scene.spheres[i])
      if (t > 0) {
        in_shadow = true
        break
      }
    }
  }

  var ambient = 0.15
  var intensity = in_shadow ? ambient : ambient + diffuse * 0.85
  return vmul(closest_sphere.color, intensity)
}

function render(width, height, scene) {
  var aspect = width / height
  var fov = 1.0
  var total_r = 0
  var total_g = 0
  var total_b = 0
  var y = 0
  var x = 0
  var u = 0
  var v = 0
  var dir = null
  var color = null
  var origin = vec(0, 0, 0)

  for (y = 0; y < height; y++) {
    for (x = 0; x < width; x++) {
      u = (2 * (x + 0.5) / width - 1) * aspect * fov
      v = (1 - 2 * (y + 0.5) / height) * fov
      dir = vnorm(vec(u, v, 1))
      color = trace(origin, dir, scene)
      total_r += color.x
      total_g += color.y
      total_b += color.z
    }
  }

  return {r: total_r, g: total_g, b: total_b}
}

var scene = make_scene()

return {
  raytrace_32x32: function(n) {
    var i = 0
    var result = null
    for (i = 0; i < n; i++) {
      result = render(32, 32, scene)
    }
    return result
  },

  raytrace_64x64: function(n) {
    var i = 0
    var result = null
    for (i = 0; i < n; i++) {
      result = render(64, 64, scene)
    }
    return result
  }
}