// The Benchmarks Game's `n-body`: the sun and four planets, moved forward in
// small time steps, with the system's total energy printed before and after.
//
//     nikaia run -- 1000
//     -0.169075164
//     -0.169087605
//
// The argument is the number of steps (1000 if left out). These two numbers
// are the benchmark's published output for 1000 steps.
//
// Shows: structs with methods, `ref mut self` for methods that change their
// value, `sync` functions, and floating-point arithmetic in a loop.
//
// A package: `nikaia run` in this directory runs it, and `nikaia test` checks
// what it prints against `tests/` (ADR-269 D13).

use std::cli
use std::text

// One body: a position, a velocity, and a mass. All of it is changed in place,
// since the step below runs millions of times.
pub struct Body {
    x: f64,
    y: f64,
    z: f64,
    vx: f64,
    vy: f64,
    vz: f64,
    mass: f64,
}

// The bodies as a set, because the interesting operations - the pairwise
// force, the total energy - are about all of them at once.
pub struct System {
    bodies: Vec[Body],
}

impl System {

    // The sun's velocity is not given: it is whatever makes the total momentum
    // of the system zero, which keeps the centre of mass still.
    //
    // `sync` says the method cannot pause (no file or network I/O, for
    // example); the compiler refuses the program if it could.
    fn offset_momentum(ref mut self) sync {
        let mut px = 0.0
        let mut py = 0.0
        let mut pz = 0.0

        for i in 0..<self.bodies.len() {
            px = px + self.bodies[i].vx * self.bodies[i].mass
            py = py + self.bodies[i].vy * self.bodies[i].mass
            pz = pz + self.bodies[i].vz * self.bodies[i].mass
        }

        let solar_mass = 4.0 * 3.141592653589793 * 3.141592653589793
        self.bodies[0].vx = -px / solar_mass
        self.bodies[0].vy = -py / solar_mass
        self.bodies[0].vz = -pz / solar_mass
    }

    // One step: first every pair pulls on both of its bodies, then every body
    // moves. Updating all velocities before any position is what keeps the
    // energy from drifting over millions of steps, so the two loops cannot be
    // merged into one.
    //
    // The inner loop starts at `i + 1`: each pair is visited once, and the pull
    // on its two bodies differs only in sign.
    fn advance(ref mut self, dt: f64) sync {
        let n = self.bodies.len()

        for i in 0..<n {
            for j in i + 1..<n {
                let dx = self.bodies[i].x - self.bodies[j].x
                let dy = self.bodies[i].y - self.bodies[j].y
                let dz = self.bodies[i].z - self.bodies[j].z

                let d2 = dx * dx + dy * dy + dz * dz
                let mag = dt / (d2 * d2.sqrt())

                let mi = self.bodies[i].mass
                let mj = self.bodies[j].mass

                self.bodies[i].vx = self.bodies[i].vx - dx * mj * mag
                self.bodies[i].vy = self.bodies[i].vy - dy * mj * mag
                self.bodies[i].vz = self.bodies[i].vz - dz * mj * mag

                self.bodies[j].vx = self.bodies[j].vx + dx * mi * mag
                self.bodies[j].vy = self.bodies[j].vy + dy * mi * mag
                self.bodies[j].vz = self.bodies[j].vz + dz * mi * mag
            }
        }

        for i in 0..<n {
            self.bodies[i].x = self.bodies[i].x + dt * self.bodies[i].vx
            self.bodies[i].y = self.bodies[i].y + dt * self.bodies[i].vy
            self.bodies[i].z = self.bodies[i].z + dt * self.bodies[i].vz
        }
    }

    // Kinetic energy of each body, less the potential energy of each pair. An
    // exact integrator would keep it constant, so printing it before and after
    // shows how well the arithmetic held up.
    fn energy(ref self) -> f64 sync {
        let n = self.bodies.len()
        let mut e = 0.0

        for i in 0..<n {
            let b = self.bodies[i]
            e = e + 0.5 * b.mass * (b.vx * b.vx + b.vy * b.vy + b.vz * b.vz)

            for j in i + 1..<n {
                let dx = b.x - self.bodies[j].x
                let dy = b.y - self.bodies[j].y
                let dz = b.z - self.bodies[j].z
                e = e - (b.mass * self.bodies[j].mass) / (dx * dx + dy * dy + dz * dz).sqrt()
            }
        }

        return e
    }
}

// The sun, Jupiter, Saturn, Uranus and Neptune, in the benchmark's units:
// astronomical units, years, and a solar mass of 4 pi squared. The velocities
// are given per day, so each is multiplied by the days in a year.
fn solar_system() -> System {
    let days_per_year = 365.24
    let solar_mass = 4.0 * 3.141592653589793 * 3.141592653589793

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

    return System { bodies: bodies }
}

fn main() {
    // `??` is for a value that may be absent: `nth` returns one. Parsing is
    // different, it can fail, so its fallback is a `catch` whose block gives
    // the value to use instead.
    let steps = cli::args().nth(1) ?? "1000"
    let n = text::parse_i32(steps) ?? 1000

    let mut system = solar_system()
    system.offset_momentum()

    println(f"{system.energy():.9}")
    for _step in 0..<n {
        system.advance(0.01)
    }
    println(f"{system.energy():.9}")
}