# 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


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

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


class Pair:
    Body b1
    Body b2

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


typedef Bodies = array_of(Body)
typedef Pairs = array_of(Pair)
typedef String_Body_di = dict_of(String, Body)


Pairs combinations(Bodies l):
    Pairs result = array(Pairs, 0)
    for Int i = 0 to len(l) - 1:
        Body x = l[i]
        for Int j = i + 1 to len(l):
            result.append(Pair(l[i], l[j]))
    return result


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

String_Body_di BODIES = {String_Body_di
    String("sun"): Body(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, SOLAR_MASS),

    String("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),

    String("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),

    String("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),

    String("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) }


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


void advance(Real dt, Int n):
    for Int i = 0 to n:
        for Pair pair in PAIRS:
            Body b1 = pair.b1
            Body b2 = pair.b2
            
            Real dx = b1.x - b2.x
            Real dy = b1.y - b2.y
            Real dz = b1.z - b2.z

            Real mag = dt * pow(dx * dx + dy * dy + dz * dz, -1.5)
            Real b1m = b1.m * mag
            Real 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 body in SYSTEM:
            body.x = body.x + dt * body.vx
            body.y = body.y + dt * body.vy
            body.z = body.z + dt * body.vz


void report_energy():
    Real e = 0.0
    for Pair pair in PAIRS:
        Body b1 = pair.b1
        Body b2 = pair.b2
        Real dx = b1.x - b2.x
        Real dy = b1.y - b2.y
        Real dz = b1.z - b2.z
        e = e - (b1.m * b2.m) / pow(dx * dx + dy * dy + dz * dz,  0.5)

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


void offset_momentum(Body ref):
    Real px = 0
    Real py = 0
    Real pz = 0

    for Body body in SYSTEM:
        Real m = body.m
        px = px - body.vx * m
        py = py - body.vy * m
        pz = pz - body.vz * m


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


void main(Int n):
    offset_momentum(BODIES[String("sun")])
    report_energy()
    advance(0.01, n)
    report_energy()

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

