// nbody.cm — N-body gravitational simulation kernel
// Pure numeric + allocation workload. Classic VM benchmark.

var math = use('math/radians')

def PI = 3.141592653589793
def SOLAR_MASS = 4 * PI * PI
def DAYS_PER_YEAR = 365.24

function make_system() {
  // Sun + 4 Jovian planets
  var sun = {x: 0, y: 0, z: 0, vx: 0, vy: 0, vz: 0, mass: SOLAR_MASS}

  var jupiter = {
    x: 4.84143144246472090,
    y: -1.16032004402742839,
    z: -0.103622044471123109,
    vx: 0.00166007664274403694 * DAYS_PER_YEAR,
    vy: 0.00769901118419740425 * DAYS_PER_YEAR,
    vz: -0.0000690460016972063023 * DAYS_PER_YEAR,
    mass: 0.000954791938424326609 * SOLAR_MASS
  }

  var saturn = {
    x: 8.34336671824457987,
    y: 4.12479856412430479,
    z: -0.403523417114321381,
    vx: -0.00276742510726862411 * DAYS_PER_YEAR,
    vy: 0.00499852801234917238 * DAYS_PER_YEAR,
    vz: 0.0000230417297573763929 * DAYS_PER_YEAR,
    mass: 0.000285885980666130812 * SOLAR_MASS
  }

  var uranus = {
    x: 12.8943695621391310,
    y: -15.1111514016986312,
    z: -0.223307578892655734,
    vx: 0.00296460137564761618 * DAYS_PER_YEAR,
    vy: 0.00237847173959480950 * DAYS_PER_YEAR,
    vz: -0.0000296589568540237556 * DAYS_PER_YEAR,
    mass: 0.0000436624404335156298 * SOLAR_MASS
  }

  var neptune = {
    x: 15.3796971148509165,
    y: -25.9193146099879641,
    z: 0.179258772950371181,
    vx: 0.00268067772490389322 * DAYS_PER_YEAR,
    vy: 0.00162824170038242295 * DAYS_PER_YEAR,
    vz: -0.0000951592254519715870 * DAYS_PER_YEAR,
    mass: 0.0000515138902046611451 * SOLAR_MASS
  }

  var bodies = [sun, jupiter, saturn, uranus, neptune]

  // Offset momentum
  var px = 0
  var py = 0
  var pz = 0
  var i = 0
  for (i = 0; i < length(bodies); i++) {
    px += bodies[i].vx * bodies[i].mass
    py += bodies[i].vy * bodies[i].mass
    pz += bodies[i].vz * bodies[i].mass
  }
  sun.vx = -px / SOLAR_MASS
  sun.vy = -py / SOLAR_MASS
  sun.vz = -pz / SOLAR_MASS

  return bodies
}

function advance(bodies, dt) {
  var n = length(bodies)
  var i = 0
  var j = 0
  var bi = null
  var bj = null
  var dx = 0
  var dy = 0
  var dz = 0
  var dist_sq = 0
  var dist = 0
  var mag = 0

  for (i = 0; i < n; i++) {
    bi = bodies[i]
    for (j = i + 1; j < n; j++) {
      bj = bodies[j]
      dx = bi.x - bj.x
      dy = bi.y - bj.y
      dz = bi.z - bj.z
      dist_sq = dx * dx + dy * dy + dz * dz
      dist = math.sqrt(dist_sq)
      mag = dt / (dist_sq * dist)

      bi.vx -= dx * bj.mass * mag
      bi.vy -= dy * bj.mass * mag
      bi.vz -= dz * bj.mass * mag
      bj.vx += dx * bi.mass * mag
      bj.vy += dy * bi.mass * mag
      bj.vz += dz * bi.mass * mag
    }
  }

  for (i = 0; i < n; i++) {
    bi = bodies[i]
    bi.x += dt * bi.vx
    bi.y += dt * bi.vy
    bi.z += dt * bi.vz
  }
}

function energy(bodies) {
  var e = 0
  var n = length(bodies)
  var i = 0
  var j = 0
  var bi = null
  var bj = null
  var dx = 0
  var dy = 0
  var dz = 0
  for (i = 0; i < n; i++) {
    bi = bodies[i]
    e += 0.5 * bi.mass * (bi.vx * bi.vx + bi.vy * bi.vy + bi.vz * bi.vz)
    for (j = i + 1; j < n; j++) {
      bj = bodies[j]
      dx = bi.x - bj.x
      dy = bi.y - bj.y
      dz = bi.z - bj.z
      e -= (bi.mass * bj.mass) / math.sqrt(dx * dx + dy * dy + dz * dz)
    }
  }
  return e
}

return {
  nbody_1k: function(n) {
    var i = 0
    var j = 0
    var bodies = null
    for (i = 0; i < n; i++) {
      bodies = make_system()
      for (j = 0; j < 1000; j++) advance(bodies, 0.01)
      energy(bodies)
    }
  },

  nbody_10k: function(n) {
    var i = 0
    var j = 0
    var bodies = null
    for (i = 0; i < n; i++) {
      bodies = make_system()
      for (j = 0; j < 10000; j++) advance(bodies, 0.01)
      energy(bodies)
    }
  }
}