n body

load core::math::sqrt;

val NBODIES: int = 5;
val SOLAR_MASS: float = 39.47841760435743;
val DAYS_PER_YEAR: float = 365.24;

struct System {
  x: []float,
  y: []float,
  z: []float,
  vx: []float,
  vy: []float,
  vz: []float,
  mass: []float,
}

apply System {

  fun offset_momentum(mut self) {
    mut px: float = 0.0;
    mut py: float = 0.0;
    mut pz: float = 0.0;

    for i := 0..NBODIES {
      px += self.vx[i] * self.mass[i];
      py += self.vy[i] * self.mass[i];
      pz += self.vz[i] * self.mass[i];
    }

    self.vx[0] = -px / SOLAR_MASS;
    self.vy[0] = -py / SOLAR_MASS;
    self.vz[0] = -pz / SOLAR_MASS;
  }

  fun energy(self) -> float {
    mut e: float = 0.0;

    for i := 0..NBODIES {
      e += 0.5 * self.mass[i] *
        (self.vx[i] * self.vx[i] +
         self.vy[i] * self.vy[i] +
         self.vz[i] * self.vz[i]);

      imu lo := i + 1;

      for j := lo..NBODIES {
        imu dx := self.x[i] - self.x[j];
        imu dy := self.y[i] - self.y[j];
        imu dz := self.z[i] - self.z[j];
        imu dist := sqrt(dx * dx + dy * dy + dz * dz);

        e -= (self.mass[i] * self.mass[j]) / dist;
      }
    }

    e
  }

  fun advance(mut self, dt: float) {
    for i := 0..NBODIES {
      imu lo := i + 1;

      for j := lo..NBODIES {
        imu dx := self.x[i] - self.x[j];
        imu dy := self.y[i] - self.y[j];
        imu dz := self.z[i] - self.z[j];
        imu d2 := dx * dx + dy * dy + dz * dz;
        imu mag := dt / (d2 * sqrt(d2));

        self.vx[i] = self.vx[i] - dx * self.mass[j] * mag;
        self.vy[i] = self.vy[i] - dy * self.mass[j] * mag;
        self.vz[i] = self.vz[i] - dz * self.mass[j] * mag;
        self.vx[j] = self.vx[j] + dx * self.mass[i] * mag;
        self.vy[j] = self.vy[j] + dy * self.mass[i] * mag;
        self.vz[j] = self.vz[j] + dz * self.mass[i] * mag;
      }
    }

    for i := 0..NBODIES {
      self.x[i] = self.x[i] + dt * self.vx[i];
      self.y[i] = self.y[i] + dt * self.vy[i];
      self.z[i] = self.z[i] + dt * self.vz[i];
    }
  }
}

fun main() {
  imu argv: []str = args();
  imu n: int = if argv.len > 0 {
    match argv[0].parse_int() {
      Option::Some(steps) => steps,
      Option::None => 1000,
    }
  } else {
    1000
  };

  mut sys: System = System {
    x = [
      0.0,
      4.84143144246472090e+00,
      8.34336671824457987e+00,
      1.28943695621391310e+01,
      1.53796971148509165e+01,
    ],
    y = [
      0.0,
      -1.16032004402742839e+00,
      4.12479856412430479e+00,
      -1.51111514016986312e+01,
      -2.59193146099879641e+01,
    ],
    z = [
      0.0,
      -1.03622044471123109e-01,
      -4.03523417114321381e-01,
      -2.23307578892655734e-01,
      1.79258772950371181e-01,
    ],
    vx = [
      0.0,
      1.66007664274403694e-03 * DAYS_PER_YEAR,
      -2.76742510726862411e-03 * DAYS_PER_YEAR,
      2.96460137564761618e-03 * DAYS_PER_YEAR,
      2.68067772490389322e-03 * DAYS_PER_YEAR,
    ],
    vy = [
      0.0,
      7.69901118419740425e-03 * DAYS_PER_YEAR,
      4.99852801234917238e-03 * DAYS_PER_YEAR,
      2.37847173959480950e-03 * DAYS_PER_YEAR,
      1.62824170038242295e-03 * DAYS_PER_YEAR,
    ],
    vz = [
      0.0,
      -6.90460016972063023e-05 * DAYS_PER_YEAR,
      2.30417297573763929e-05 * DAYS_PER_YEAR,
      -2.96589568540237556e-05 * DAYS_PER_YEAR,
      -9.51592254519715870e-05 * DAYS_PER_YEAR,
    ],
    mass = [
      SOLAR_MASS,
      9.54791938424326609e-04 * SOLAR_MASS,
      2.85885980666130812e-04 * SOLAR_MASS,
      4.36624404335156298e-05 * SOLAR_MASS,
      5.15138902046611451e-05 * SOLAR_MASS,
    ],
  };

  sys.offset_momentum();
  showln(sys.energy());

  for _ := 0..n {
    sys.advance(0.01);
  }

  showln(sys.energy());
}