import csv
import numpy as np
from math import exp, sqrt, sin, cos, pi

"""
  Planet simulator: Solve simulataneous second order diffrential equations by Euler method
"""

# d2x/dt2 = force(t, x, y), d2y/dt2 = force(t, x, y)
# define function to be integrated
def force(t, x, y):
    return -x, -y

# solution: x(0.0) = 0.0, y(0.0) = 2.0, vx(0.0) = 1.0, vy(0.0) = 0.0
def xysolution(t):
    return sin(t), 2.0 * cos(t)

#===================
# constants
#===================
G = 6.67259e-11                       #Nm2/kg2 
DayToSecond      = 60 * 60 * 24       #s
SecondToDay      = 1.0 / DayToSecond
AstronomicalUnit = 1.49597870e11      #m
AU               = AstronomicalUnit
G1 = G * DayToSecond * DayToSecond / AU / AU / AU

#===================
# parameters
#===================
dbfile  = 'planet_db.csv'
outfile = 'diffeq2nd_Planet_Euler.csv'
x0, vx0 = 0.0, 1.0
y0, vy0 = 2.0, 0.0
z0, vz0 = 0.0, 0.0
dt = 1.0
nt = 1000
iprint_interval = 100
nprint_planets = 4
CSVOutputInterval = 10

CriteriaR2ForIdenticalAtom = 0.001
DifferentialEquationSolver = "Verlet"

#===================
# functions
#===================
def readdb(dbfile):
    planets = []
    f = open(dbfile, "r");
    reader = csv.DictReader(f)
    for row in reader:
       planets.append(row)
    keys = list(planets[0].keys())
    for d in planets:
        for key in keys:
            if key != 'Name':
                d[key] = float(d[key])
    return planets

def sum(array):
    sum = 0.0
    for e in array:
        sum += e
    return sum

def total_momentum(it, M, vx, vy, vz):
    Px, Py, Pz = 0.0, 0.0, 0.0
    for i in range(0, len(M)):
        Px += M[i] * vx[i][it];
        Py += M[i] * vy[i][it];
        Pz += M[i] * vz[i][it];
    return Px, Py, Pz

# Normalize total momentum to zero
def normalize_momentum(it, M, x, y, z, vx, vy, vz, fx, fy, fz):
    Mtot = sum(M)
    Px, Py, Pz = total_momentum(it, M, vx, vy, vz);
    for ip in range(0, len(M)):
        vx[ip][it] -= Px / Mtot;
        vy[ip][it] -= Py / Mtot;
        vz[ip][it] -= Pz / Mtot;
    return Px, Py, Pz

def initialize(planets, M, x, y, z, vx, vy, vz, fx, fy, fz):
    global AU
    global DayToSecond

    for i in range(0, len(planets)):
        M.append(planets[i]['Mass'])
        x.append([planets[i]['Revolution Radius'] / AU])
        y.append([0.0])
        z.append([0.0])
        vx.append([0.0])
        vy.append([planets[i]['Revolution Velocity'] * DayToSecond / AU])
        vz.append([0.0])

    for i in range(0, len(planets)):
        fxi, fyi, fzi = Fi(0, i, M, x, y, z)
        fx.append([fxi])
        fy.append([fyi])
        fz.append([fzi])

def Fij(istep, i, j, M, x, y, z):
    dx = x[j][istep] - x[i][istep]
    dy = y[j][istep] - y[i][istep]
    dz = z[j][istep] - z[i][istep]
    r2 = dx*dx + dy*dy + dz*dz
    r  = sqrt(r2)
    g = G1 * M[j]
    f = g / r2
    fx = f * dx / r
    fy = f * dy / r
    fz = f * dz / r
    return fx, fy, fz

def Fi(istep, i, M, x, y, z):
    fxi = 0.0
    fyi = 0.0
    fzi = 0.0
    for j in range(0, len(M)):
        if i == j:
            continue
        fxj, fyj, fzj = Fij(istep, i, j, M, x, y, z)
        fxi += fxj
        fyi += fyj
        fzi += fzj
#    print("f={}, {}, {}".format(fxi, fyi, fzi))
    return fxi, fyi, fzi

#===================
# main routine
#===================
def main():
    global nt
    global dt

    print("Planet simulator: Solve simulataneous second order diffrential equations by Euler method")
    print("G = {} Nm2/kg2".format(G))
    print("AU = {:e} m".format(AU))
    print("G1 = {}".format(G1))
    print("")

    print("Planets:")
    planets = readdb(dbfile)
    keys = list(planets[0].keys())
    for d in planets:
        print("  ", d['Name'])
        for key in keys:
            if key != 'Name':
                print("     {}: {}".format(key, d[key]))
    print("")

    M  = []
    x  = []
    y  = []
    z  = []
    vx = []
    vy = []
    vz = []
    fx  = []
    fy  = []
    fz  = []
    initialize(planets, M, x, y, z, vx, vy, vz, fx, fy, fz)
    Px, Py, Pz = normalize_momentum(0, M, x, y, z, vx, vy, vz, fx, fy, fz)
    print("")

# make label list for display / csv output
    labellist = ['t']
    for i in range(0, len(planets)):
        labellist.append("x({})".format(planets[i]['Name']))
        labellist.append("y({})".format(planets[i]['Name']))

# open outfile to write a csv file
    print("Write to [{}]".format(outfile))
    f = open(outfile, 'w')
    fout = csv.writer(f, lineterminator='\n')
    fout.writerow(labellist)

    print("{:^5}".format('t'), end = '')
    for i in range(1, nprint_planets*2, 2):
        print("  {:^12}  {:^12}".format(labellist[i], labellist[i+1]), end = '')
    print("")

# Solve the 1st data by Heun method
    datalist = [0.0]
    print("{:^5}".format(0.0), end = '')
    for i in range(0, len(planets)):
        fx0, fy0, fz0 = Fi(0, i, M, x, y, z)
        vx1 = vx[i][0] + dt * fx0
        vy1 = vy[i][0] + dt * fy0
        vz1 = vz[i][0] + dt * fz0
        x1  = x[i][0]  + dt * vx[i][0]
        y1  = y[i][0]  + dt * vy[i][0]
        z1  = z[i][0]  + dt * vz[i][0]
        datalist.append(x[i][0])
        datalist.append(y[i][0])
        x[i].append(x1)
        y[i].append(y1)
        z[i].append(z1)
        vx[i].append(vx1)
        vy[i].append(vy1)
        vz[i].append(vz1)
    for i in range(1, nprint_planets*2, 2):
        print("  {:^12.4f}  {:^12.4f}".format(x[i][0], y[i][0]), end = '')
    print("")
    fout.writerow(datalist)

    for it in range(1, nt):
        t = it * dt
#        print("it={}  t={}".format(it, t))
        datalist = [t]
        if it % iprint_interval == 0:
            print("{:^5}".format(t), end = '')
        for i in range(0, len(planets)):
            fx0, fy0, fz0 = Fi(it, i, M, x, y, z)
            vx1 = vx[i][it] + dt * fx0
            vy1 = vy[i][it] + dt * fy0
            vz1 = vz[i][it] + dt * fz0
            x1  = x[i][it]  + dt * vx[i][it]
            y1  = y[i][it]  + dt * vy[i][it]
            z1  = z[i][it]  + dt * vz[i][it]
            datalist.append(x[i][it])
            datalist.append(y[i][it])
            x[i].append(x1)
            y[i].append(y1)
            z[i].append(z1)
            vx[i].append(vx1)
            vy[i].append(vy1)
            vz[i].append(vz1)
        if it % iprint_interval == 0:
            for i in range(1, nprint_planets*2, 2):
                print("  {:^12.4g}  {:^12.4g}".format(x[i][it], y[i][it]), end = '')
            print("")
        fout.writerow(datalist)

    f.close()


if __name__ == '__main__':
    main()

