import math
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
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"]
xposs = []
yposs = []
zposs = []

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
    xpos = data["X"].mean()
    ypos = data["Y"].mean()
    zpos = data["Z"].mean()
    xposs.append(xpos)
    yposs.append(ypos)
    zposs.append(zpos)

theta = np.linspace(0, 2 * np.pi, 100)
phi = np.linspace(0, np.pi, 50)
theta, phi = np.meshgrid(theta, phi)
r = 69911

x = r * np.sin(phi) * np.cos(theta)
y = r * np.sin(phi) * np.sin(theta)
z = r * np.cos(phi)


ax=plt.axes(projection='3d')
ax.set_box_aspect([1, 1, 1])
max_range = 0.5*np.max([np.max(np.abs(xposs)), np.max(np.abs(yposs)), np.max(np.abs(zposs))])
ax.set_xlim(-max_range, max_range)
ax.set_ylim(-max_range, max_range)
ax.set_zlim(-max_range, max_range)
ax.plot(xposs, yposs, zposs)
ax.plot_surface(x, y, z, alpha=0.3)
plt.show()