# The Computer Language Benchmarks Game
# https://salsa.debian.org/benchmarksgame-team/benchmarksgame/
#
# originally by Kevin Carson
# modified by Tupteq, Fredrik Johansson, and Daniel Nanz
# modified by Maciej Fijalkowski
# 2to3

require "debug"

def combinations(l):
    result = []
    for i = 0 to len(l) - 1:
        var x = l[i]
        for j = i + 1 to len(l):
            result.append(Pair(l[i], l[j]))
    return result


class Pair:
    var b1
    var b2

    def init(b1_, b2_):
        b1 = b1_
        b2 = b2_


class Body:
    var x, y, z
    var vx, vy, vz
    var m

    def init(x_, y_, z_, vx_, vy_, vz_, m_):
        x = x_
        y = y_
        z = z_
        vx = vx_
        vy = vy_
        vz = vz_
        m = m_


PI = 3.14159265358979323
SOLAR_MASS = 4 * PI * PI
DAYS_PER_YEAR = 365.24

BODIES = {
    'sun': Body(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, SOLAR_MASS),

    'jupiter': Body(4.84143144246472090e+00,
                    -1.16032004402742839e+00,
                    -1.03622044471123109e-01,
                    1.66007664274403694e-03 * DAYS_PER_YEAR,
                    7.69901118419740425e-03 * DAYS_PER_YEAR,
                    -6.90460016972063023e-05 * DAYS_PER_YEAR,
                    9.54791938424326609e-04 * SOLAR_MASS),

    'saturn': Body(8.34336671824457987e+00,
                   4.12479856412430479e+00,
                   -4.03523417114321381e-01,
                   -2.76742510726862411e-03 * DAYS_PER_YEAR,
                   4.99852801234917238e-03 * DAYS_PER_YEAR,
                   2.30417297573763929e-05 * DAYS_PER_YEAR,
                   2.85885980666130812e-04 * SOLAR_MASS),

    'uranus': Body(1.28943695621391310e+01,
                   -1.51111514016986312e+01,
                   -2.23307578892655734e-01,
                   2.96460137564761618e-03 * DAYS_PER_YEAR,
                   2.37847173959480950e-03 * DAYS_PER_YEAR,
                   -2.96589568540237556e-05 * DAYS_PER_YEAR,
                   4.36624404335156298e-05 * SOLAR_MASS),

    'neptune': Body(1.53796971148509165e+01,
                    -2.59193146099879641e+01,
                    1.79258772950371181e-01,
                    2.68067772490389322e-03 * DAYS_PER_YEAR,
                    1.62824170038242295e-03 * DAYS_PER_YEAR,
                    -9.51592254519715870e-05 * DAYS_PER_YEAR,
                    5.15138902046611451e-05 * SOLAR_MASS) }


SYSTEM = BODIES.values()
PAIRS = combinations(SYSTEM)


def advance(dt, n, keyword bodies=SYSTEM, pairs=PAIRS):
    for i = 0 to n:
        for pair in pairs:
            var b1 = pair.b1
            var b2 = pair.b2
            
            var dx = b1.x - b2.x
            var dy = b1.y - b2.y
            var dz = b1.z - b2.z

            var mag = dt * ((dx * dx + dy * dy + dz * dz) ** (-1.5))
            var b1m = b1.m * mag
            var b2m = b2.m * mag
            b1.vx = b1.vx - dx * b2m
            b1.vy = b1.vy - dy * b2m
            b1.vz = b1.vz - dz * b2m
            b2.vx = b2.vx + dx * b1m
            b2.vy = b2.vy + dy * b1m
            b2.vz = b2.vz + dz * b1m

        for body in bodies:
            body.x = body.x + dt * body.vx
            body.y = body.y + dt * body.vy
            body.z = body.z + dt * body.vz


def report_energy(keyword bodies=SYSTEM, pairs=PAIRS, e=0.0):
    for pair in pairs:
        var b1 = pair.b1
        var b2 = pair.b2
        var dx = b1.x - b2.x
        var dy = b1.y - b2.y
        var dz = b1.z - b2.z
        e = e - (b1.m * b2.m) / ((dx * dx + dy * dy + dz * dz) ** 0.5)

    for body in bodies:
        e = e + body.m * (body.vx * body.vx + body.vy * body.vy +
                          body.vz * body.vz) / 2.0
    print e


def offset_momentum(ref, keyword bodies=SYSTEM):
    var px = 0
    var py = 0
    var pz = 0

    for body in bodies:
        var m = body.m
        px = px - body.vx * m
        py = py - body.vy * m
        pz = pz - body.vz * m


    m = ref.m
    ref.vx = px / m
    ref.vy = py / m
    ref.vz = pz / m


def main(n, keyword ref='sun'):
    offset_momentum(BODIES[ref])
    report_energy()
    advance(0.01, n)
    report_energy()

print "nbody", command_line_arguments
main(int(command_line_arguments[1]))
exit()

