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

def kepler_to_cartesian(a, e, i, Omega, w, nu, G=1.0, M=1.0):
    def to_arrays(*args):
        result = []
        for arg in args: result.append(np.atleast_1d(arg))
        return result
    a, e, i, Omega, w, nu = to_arrays(a, e, i, Omega, w, nu) # pylint: disable=unbalanced-tuple-unpacking
    P = [np.cos(w)*np.cos(Omega) - np.sin(w)*np.cos(i)*np.sin(Omega),
         np.cos(w)*np.sin(Omega) + np.sin(w)*np.cos(i)*np.cos(Omega),
         np.sin(w)*np.sin(i)]
    Q = [-np.sin(w)*np.cos(Omega) - np.cos(w)*np.cos(i)*np.sin(Omega),
         -np.sin(w)*np.sin(Omega) + np.cos(w)*np.cos(i)*np.cos(Omega),
          np.sin(i)*np.cos(w)]
    cosnu = np.cos(nu)
    cosE = (e+cosnu)/(1+e*cosnu)
    E = np.arccos(cosE)
    E[nu > np.pi] = 2*np.pi - E[nu > np.pi]
    X = a*((np.cos(E)-e)*P + np.sqrt(1-e**2)*np.sin(E)*Q)
    V = (np.sqrt(G*M/a)/(1-e*np.cos(E)))*(-np.sin(E)*P + np.sqrt(1-e**2)*np.cos(E)*Q)
    if X.shape[1]==1:
        X=X[:,0]
        V=V[:,0]
    return X.T, V.T

def generate_binary(a, e, i, Omega, w, nu, m1, m2):
    X, V = kepler_to_cartesian(a, e, i, Omega, w, nu, M=m1+m2)
    q = np.double(m2)/np.double(m1)
    X1 = -q/(q+1)*X
    V1 = -q/(q+1)*V
    X2 =  1/(q+1)*X
    V2 =  1/(q+1)*V
    return X1, V1, X2, V2

def write_phi_grape_config(**kargs):
    if 'file_name' in kargs: file_name = kargs['file_name']
    else: file_name = 'phi-GRAPE.cfg'
    f = open(file_name, 'w')
    f.write('%.4E %.4E %.4E %.4E %.4E %.4E data.con\n' % (kargs['eps'], kargs['t_end'], kargs['dt_disk'], kargs['dt_contr'], kargs['dt_bh'], kargs['eta']))
    f.close()

def write_particles(particle_list, m=None, file_name='data.con'):
    N = particle_list.shape[0]
    if m is None:
        m = np.empty(N)
        m[:] = 1/N
    f = open(file_name, '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()

def gen_mask(particle_list, frac):
    N = particle_list.shape[0]
    if   frac==0:
        mask = np.ones(N, dtype=int)
    elif frac==1:
        mask = np.zeros(N, dtype=int)
    elif (frac < 0) or (1 < frac):
        raise RuntimeError('Fraction has to be between 0 and 1')
    else:
        X = particle_list[:,:3]
        V = particle_list[:,3:]
        L = np.linalg.norm(np.cross(X, V), axis=1)
        L_sorted = L.copy()
        L_sorted.sort()
        L_cutoff = L_sorted[int(N*frac)]
        mask = np.array(L > L_cutoff, dtype=int)
    return mask

def write_mask(mask, file_name='grapite.mask'):
    f = open(file_name, 'w')
    for i in range(len(mask)):
        f.write('%06d %d\n' % (i, mask[i]))
    f.close()

if __name__=='__main__':
    parser = argparse.ArgumentParser(description='Process some integers.')
    parser.add_argument('N',          type=str,                     help='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)')
    parser.add_argument('--frac',     type=np.double, default=0,    help='fraction of collisional particles (by angular momentum)')
    parser.add_argument('--bsmbh',    type=bool,      default=0,    help='generate a binary supermassive black hole (parameters hardcoded in the script)')
    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.')

    write_phi_grape_config(**vars(args))

    particle_list = gen_plum(N, seed=args.seed)
    m = np.ones(N)/N

    if args.bsmbh:
        m1, m2 = 0.075, 0.025
        a, e, i, Omega, w, nu = 0.001, 0.5, 0, 0, 0, 0
        X1, V1, X2, V2 = generate_binary(a, e, i, Omega, w, nu, m1, m2)
        m[:2]  = m1, m2
        m[2:] = 1/(N-2)
        particle_list[0,:3] = X1
        particle_list[0,3:] = V1
        particle_list[1,:3] = X2
        particle_list[1,3:] = V2

    write_particles(particle_list, m=m)

    mask = gen_mask(particle_list, args.frac)
    if args.bsmbh:
        mask[:2] = 3

    write_mask(mask)