import math
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--density", "-d")

args=parser.parse_args()

columns = ["ID", "idx", "Mass", "Radius", "X", "Y", "Z", "vX", "vY", "vZ", "sX", "sY", "sZ", "Colour"]
Kenergies = []
Genergies = []
energies = []
time = []

for i in range (1,401):
    num = str(i).rjust(5, '0')
    file = "BTs/"+str(args.density)+"/boom."+num+".bt"
    data = pd.read_csv(file, sep=' ', names=columns)
    data["X"] = data["X"]*1.5e8
    data["Y"] = data["Y"]*1.5e8
    data["Z"] = data["Z"]*1.5e8
    data["X2"] = data["X"]**2
    data["Y2"] = data["Y"]**2
    data["Z2"] = data["Z"]**2
    data["distance2"] = data["X2"]+data["Y2"]+data["Z2"]
    data["distance"] = data["distance2"]**0.5
    data["vX"] = data["vX"]*29880
    data["vY"] = data["vY"]*29880
    data["vZ"] = data["vZ"]*29880
    data["vX2"] = data["vX"]**2
    data["vY2"] = data["vY"]**2
    data["vZ2"] = data["vZ"]**2
    data["Mass"] = data["Mass"]*1.989e30
    data["vel2"] = data["X"]+data["Y"]+data["Z"]**2
    data["vel"] = data["distance2"]**0.5
    data["KE"] = 0.5*data["Mass"]*data["vel2"]
    data["GPE"] = (6.67e-11*1.89e27*data["Mass"])/(data["distance"])
    data["Energy"] = data["KE"]+data["GPE"]
    energy = data["Energy"].sum()
    energies.append(energy)
    time.append(i)
    KE = data["KE"].sum()
    Kenergies.append(KE)
    GPE = data["GPE"].sum()
    Genergies.append(GPE)

ax=plt.axes()
# ax.set_yscale("log")
ax.plot(time, Kenergies, label="Kinetic")
ax.plot(time, Genergies, label="Gravitational")
ax.plot(time, energies, label="Total")
ax.legend()
plt.show()