import numpy as np
import argparse

def gen_plum(N, seed=None, RMAX=10):
    if not seed is None: np.random.seed(seed)
    particle_list = np.zeros((N,6))
    i = 0
    while (i < N):
        X1, X2, X3 = np.random.random(3)
        R = 1/np.sqrt(X1**(-2/3) - 1)
        if (R > RMAX): continue
        Z = (1 - 2*X2)*R
        X = np.sqrt(R**2 - Z**2) * np.cos(2*np.pi*X3)
        Y = np.sqrt(R**2 - Z**2) * np.sin(2*np.pi*X3)

        Ve = np.sqrt(2)*(1.0 + R**2)**(-0.25);

        X4, X5 = 0, 0
        while 0.1*X5 >= X4**2*(1-X4**2)**3.5:
            X4, X5 = np.random.random(2)

        V = Ve*X4
        X6, X7 = np.random.random(2)

        Vz = (1 - 2*X6)*V;
        Vx = np.sqrt(V**2 - Vz**2) * np.cos(2*np.pi*X7);
        Vy = np.sqrt(V**2 - Vz**2) * np.sin(2*np.pi*X7);

        X, Y, Z    = np.array([X,  Y,  Z])*3*np.pi/16
        Vx, Vy, Vz = np.array([Vx, Vy, Vz])/np.sqrt(3*np.pi/16)

        particle_list[i,:] = [X, Y, Z, Vx, Vy, Vz]
        i += 1
    return particle_list

parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('N',          type=str,                   help='The number of particles (follow by k to multiply by 1024)')
parser.add_argument('--seed',     type=int,       default=42,   help='Random seed')
parser.add_argument('--eps',      type=np.double, default=1E-4, help='Plummer softening parameter (can be even 0)')
parser.add_argument('--t_end',    type=np.double, default=4,    help='end time of calculation')
parser.add_argument('--dt_disk',  type=np.double, default=1,    help='interval of snapshot files output (0xxx.dat)')
parser.add_argument('--dt_contr', type=np.double, default=.125, help='interval for the energy control output (contr.dat)')
parser.add_argument('--dt_bh',    type=np.double, default=.125, help='interval for BH output (bh.dat & bh_nb.dat)')
parser.add_argument('--eta',      type=np.double, default=.01,  help='parameter for timestep determination (0.02 or 0.01)')


args = parser.parse_args()
try:
    N = int(args.N)
except ValueError:
    if args.N[-1]=='k': N = int(args.N[:-1])*1024
    else: raise ValueError('Unable to parse N.')

f = open('phi-GRAPE.cfg', 'w')
f.write('%.4E %.4E %.4E %.4E %.4E %.4E data.con\n' % (args.eps, args.t_end, args.dt_disk, args.dt_contr, args.dt_bh, args.eta))
f.close()


particle_list = gen_plum(N)
m = np.empty(N)
m[:] = 1/N

f = open('data.con', 'w')
f.write('000000\n')
f.write('%d\n' % N)
f.write('0.0000000000E+00\n')
for i in range(N):
    f.write('%06d%14.6e%14.6e%14.6e%14.6e%14.6e%14.6e%14.6e\n' % (i, m[i], particle_list[i,0], particle_list[i,1], particle_list[i,2], particle_list[i,3], particle_list[i,4], particle_list[i,5]))
f.close()

f = open('grapite.mask', 'w')
for i in range(N):
    f.write('%06d %d\n' % (i, 0))
f.close()