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"]
distances = []
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)
    xPosAU = data["X"].mean()
    xPos = xPosAU*1.5e8
    yPosAU = data["Y"].mean()
    yPos = xPosAU*1.5e8
    zPosAU = data["Z"].mean()
    zPos = xPosAU*1.5e8
    distance = math.sqrt((xPos**2)+(yPos**2)+(zPos**2))
    time.append(i)
    distances.append(distance)

print(min(distances))

ax=plt.axes()
# ax.set_yscale("log")
ax.plot(time, distances)
plt.show()