const PI = 3.141592653589793;
const SOLAR_MASS = 4.0 * PI * PI;
const DAYS_PER_YEAR = 365.24;

fn advance(bodies, dt) {
    for i in 0..bodies.len() {
        let bi = bodies[i];
        let bix = bi.x;
        let biy = bi.y;
        let biz = bi.z;
        let bimass = bi.mass;
        let bivx = bi.vx;
        let bivy = bi.vy;
        let bivz = bi.vz;

        for j in i + 1..bodies.len() {
            let bj = bodies[j];
            let dx = bix - bj.x;
            let dy = biy - bj.y;
            let dz = biz - bj.z;
            let mag = (dx * dx + dy * dy + dz * dz).sqrt();
            mag = dt / (mag * mag * mag);
            let 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);
        }

        bi.vx = bivx;
        bi.vy = bivy;
        bi.vz = bivz;
        bi.x = bix + dt * bivx;
        bi.y = biy + dt * bivy;
        bi.z = biz + dt * bivz;
    }
}

fn energy(bodies) {
    let e = 0.0;

    for i in 0..bodies.len() {
        let bi = bodies[i];
        let vx = bi.vx;
        let vy = bi.vy;
        let vz = bi.vz;
        let bim = bi.mass;

        e = e + (0.5 * bim * (vx * vx + vy * vy + vz * vz));

        for j in i + 1..bodies.len() {
            let bj = bodies[j];
            let dx = bi.x - bj.x;
            let dy = bi.y - bj.y;
            let dz = bi.z - bj.z;
            let distance = (dx * dx + dy * dy + dz * dz).sqrt();
            e = e - ((bim * bj.mass) / distance);
        }
    }

    e
}

fn offset_momentum(bodies) {
    let px = 0.0;
    let py = 0.0;
    let pz = 0.0;

    for i in 0..bodies.len() {
        let bi = bodies[i];
        let bim = bi.mass;
        px = px + (bi.vx * bim);
        py = py + (bi.vy * bim);
        pz = pz + (bi.vz * bim);
    }

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

fn bodies() {
    [
        // Sun
        #{
            x: 0.0,
            y: 0.0,
            z: 0.0,
            vx: 0.0,
            vy: 0.0,
            vz: 0.0,
            mass: SOLAR_MASS,
        },
        // Jupiter
        #{
            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,
        },
        // Saturn
        #{
            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,
        },
        // Uranus
        #{
            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,
        },
        // Neptune
        #{
            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,
        },
    ]
}

#[test]
pub fn nbodies_validate() {
    const EXPECTED = -0.16908926275527306;
    const N = 20000;

    let bodies = bodies();

    offset_momentum(bodies);

    for i in 0..N {
        advance(bodies, 0.01);
    }

    let e = energy(bodies);
    assert!((e - EXPECTED).abs() < 0.0000001, "{e}");
}

#[bench]
pub fn nbodies(b) {
    const N = 20000;

    let bodies = bodies();

    b.iter(|| {
        offset_momentum(bodies);

        for i in 0..N {
            advance(bodies, 0.01);
        }

        energy(bodies)
    });
}
