import numpy as np
import random
import sys
import math

G = 4.300918e-03  # pc^3 Myr^-2 M_sun^-1
MASS = 1e10  # Msun
RADIUS = 2.0  # kpc

##############################################################

filename = sys.argv[1]
N = int(sys.argv[2])

##############################################################

rcut = 10 * RADIUS

print("File:", filename)
print("Npart:", N)

# Initialise the particle properties
mass = MASS / N
pos = np.zeros((N, 3))
vel = np.zeros((N, 3))

com = np.zeros(3)
cov = np.zeros(3)

# Generate the particles
for i in range(N):

    ri = 2.0 * rcut
    while ri > rcut:
        x1 = random.random()
        ri = 1.0 / math.sqrt(x1 ** (-2.0 / 3.0) - 1.0)

    phi = 2.0 * math.pi * random.random()
    z = (1.0 - 2.0 * random.random()) * ri
    R = math.sqrt(ri * ri - z * z)

    pos[i, 0] = R * math.cos(phi)
    pos[i, 1] = R * math.sin(phi)
    pos[i, 2] = z

    # Escape velocity
    ve = math.sqrt(2.0 / math.sqrt(1.0 + ri * ri))

    t1 = 1
    t2 = 0
    while t1 > t2:
        x2 = random.random()
        t1 = random.random() * 0.1
        t2 = x2 * x2 * (1.0 - x2 * x2) ** 3.5

    vr = ve * x2

    phi = 2.0 * math.pi * random.random()
    vz = (1.0 - 2.0 * random.random()) * vr
    vR = math.sqrt(vr * vr - vz * vz)

    vel[i, 0] = vR * math.cos(phi)
    vel[i, 1] = vR * math.sin(phi)
    vel[i, 2] = vz

    # Centre of mass and velocity
    com[0] += pos[i, 0]
    com[1] += pos[i, 1]
    com[2] += pos[i, 2]
    cov[0] += vel[i, 0]
    cov[1] += vel[i, 1]
    cov[2] += vel[i, 2]


# Normalize
com /= N
cov /= N

# Scale to analytical expectation
sx = RADIUS
sv = (
    (1.02249e-6 / 1000.0)
    / sx
    * math.sqrt((RADIUS * 1000.0) ** 3.0)
    / (2.0 * 2.2489e-15 * MASS)
)

pos[:, 0] -= com[0]
pos[:, 1] -= com[1]
pos[:, 2] -= com[2]
pos *= sx * 1000.0  # Convert to pc

vel[:, 0] -= cov[0]
vel[:, 1] -= cov[1]
vel[:, 2] -= cov[2]
vel /= sv / 1.02273  # Convert to pc / Myr