// based on https://benchmarksgame-team.pages.debian.net/benchmarksgame/program/nbody-lua-4.html

var sun = @{};
var jupiter = @{};
var saturn = @{};
var uranus = @{};
var neptune = @{};

{
  var sqrt = proc(x) pow(x, 0.5);
  var e = proc(x) pow(10, x);
  var pi = 3.141592653589793;
  var solarMass = 4 * pi * pi;
  var daysPerYear = 365.24;
  sun.x = 0.0;
  sun.y = 0.0;
  sun.z = 0.0;
  sun.vx = 0.0;
  sun.vy = 0.0;
  sun.vz = 0.0;
  sun.mass = solarMass;
  jupiter.x = 4.84143144246472090;
  jupiter.y = -1.16032004402742839;
  jupiter.z = -1.03622044471123109 * e(-1);
  jupiter.vx = 1.66007664274403694 * e(-3) * daysPerYear;
  jupiter.vy = 7.69901118419740425 * e(-3) * daysPerYear;
  jupiter.vz = -6.90460016972063023 * e(-5) * daysPerYear;
  jupiter.mass = 9.54791938424326609 * e(-4) * solarMass;
  saturn.x = 8.34336671824457987;
  saturn.y = 4.12479856412430479;
  saturn.z = -4.03523417114321381 * e(-1);
  saturn.vx = -2.76742510726862411 * e(-3) * daysPerYear;
  saturn.vy = 4.99852801234917238 * e(-03) * daysPerYear;
  saturn.vz = 2.30417297573763929 * e(-05) * daysPerYear;
  saturn.mass = 2.85885980666130812 * e(-04) * solarMass;
  uranus.x = 1.28943695621391310 * e(01);
  uranus.y = -1.51111514016986312 * e(01);
  uranus.z = -2.23307578892655734 * e(-01);
  uranus.vx = 2.96460137564761618 * e(-03) * daysPerYear;
  uranus.vy = 2.37847173959480950 * e(-03) * daysPerYear;
  uranus.vz = -2.96589568540237556 * e(-05) * daysPerYear;
  uranus.mass = 4.36624404335156298 * e(-05) * solarMass;
  neptune.x = 1.53796971148509165 * e(01);
  neptune.y = -2.59193146099879641 * e(01);
  neptune.z = 1.79258772950371181 * e(-01);
  neptune.vx = 2.68067772490389322 * e(-03) * daysPerYear;
  neptune.vy = 1.62824170038242295 * e(-03) * daysPerYear;
  neptune.vz = -9.51592254519715870 * e(-05) * daysPerYear;
  neptune.mass = 5.15138902046611451 * e(-05) * solarMass;

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

  var advance = proc(bodies, nbody, dt) {
    var i = 0;
    while (i < nbody) {
      var bi = bodies[i];
      var bix = bi.x;
      var biy = bi.y;
      var biz = bi.z;
      var bimass = bi.mass;
      var bivx = bi.vx;
      var bivy = bi.vy;
      var bivz = bi.vz;
      var j = i + 1;
      while (j < nbody) {
        var bj = bodies[j];
        var dx = bix - bj.x;
        var dy = biy - bj.y;
        var dz = biz - bj.z;
        var dist2 = dx * dx + dy * dy + dz * dz;
        var mag = sqrt(dist2);
        mag = dt / (mag * dist2);
        var bm = bj.mass * mag;
        bivx = bivx - (dx * bm);
        bivy = bivy - (dy * bm);
        bivz = bivz - (dz * bm);
        bm = bimass * mag;
        bj.vx = bj.vx + (dx * bm);
        bj.vy = bj.vy + (dy * bm);
        bj.vz = bj.vz + (dz * bm);
        j = j + 1;
      };
      bi.vx = bivx;
      bi.vy = bivy;
      bi.vz = bivz;
      bi.x = bix + dt * bivx;
      bi.y = biy + dt * bivy;
      bi.z = biz + dt * bivz;
      i = i + 1;
    };
  };

  var energy = proc(bodies, nbody) {
    var e = 0;
    var i = 0;
    while (i < nbody) {
      var bi = bodies[i];
      var vx = bi.vx;
      var vy = bi.vy;
      var vz = bi.vz;
      var bim = bi.mass;
      e = e + (0.5 * bim * (vx * vx + vy * vy + vz * vz));
      var j = i + 1;
      while (j < nbody) {
        var bj = bodies[j];
        var dx = bi.x - bj.x;
        var dy = bi.y - bj.y;
        var dz = bi.z - bj.z;
        var distance = sqrt(dx * dx + dy * dy + dz * dz);
        e = e - ((bim * bj.mass) / distance);
        j = j + 1;
      };
      i = i + 1;
    };
    e
  };

  var offsetMomentum = proc(b, nbody) {
    var px = 0;
    var py = 0;
    var pz = 0;
    var i = 0;
    while(i < nbody) {
      var bi = b[i];
      var bim = bi.mass;
      px = px + (bi.vx * bim);
      py = py + (bi.vy * bim);
      pz = pz + (bi.vz * bim);
      i = i + 1;
    };
    b[0].vx = -px / solarMass;
    b[0].vy = -py / solarMass;
    b[0].vz = -pz / solarMass;
  };


  // main
  var N = 1000000;
  var nbody = #bodies;

  offsetMomentum(bodies, nbody);
  print(energy(bodies, nbody));
  
  var i = 1;
  while (i <= N) {
    advance(bodies, nbody, 0.01);
    i = i + 1;
  };

  print(energy(bodies, nbody));
};