#!/usr/bin/env python
import sys
import numpy as np, pandas as pd

import importlib.util
spec = importlib.util.spec_from_file_location('scf', '../etics/src/scf.py')
scf = importlib.util.module_from_spec(spec)
spec.loader.exec_module(scf)

n_max        = 15
l_max        = 2
length_scale = 10

pos_filename   = '../../../sobolenko/TNG50-copy/for_yohai/98990136-weight.dat'
grav_filename  = 'out.dat'
coeff_filename = 'out.coeff'

pos_data = np.genfromtxt(pos_filename, skip_header=1, names=['idx', 'm', 'x', 'y', 'z', 'vx', 'vy', 'vz', 'eps', 'type', 'name', 't'], dtype=[int, float, float, float, float, float, float, float, float, int, int, int])
grav_data = np.genfromtxt(grav_filename, names=['name', 'pot', 'ax', 'ay', 'az'], dtype=[int, float, float, float, float])

coeff_count = (n_max+1)*(l_max+1)*(l_max+2)//2
A = np.fromfile(coeff_filename, dtype=np.complex128)[:coeff_count]
expansion = scf.Scf(n_max, l_max, length_scale=length_scale)
expansion.set_coefficients(A)

r = np.sqrt(pos_data['x']**2 + pos_data['y']**2 + pos_data['z']**2)
r[r==0] = 1 # Will be ignored later
density_all = expansion.calc_radial_density(r)

n = len(pos_data['idx'])
i2 = 0
out_file = open('results.dat', 'w')
for i1 in range(10):
    template = '%08d%15.8e%16.8e%16.8e%16.8e%16.8e%16.8e%16.8e%11.4e %d %015d%5.04d%16.8e%16.8e%16.8e%16.8e%15.8e\n'
    if pos_data['name'][i1] == grav_data['name'][i2]:
        pot, ax, ay, az = grav_data['pot'][i2], grav_data['ax'][i2], grav_data['ay'][i2], grav_data['az'][i2]
        den = density_all[i1]
        i2 += 1
    else: # Missing particle
        pot, ax, ay, az = 0, 0, 0, 0
        den = 0
    out = template % (pos_data['idx'][i1],
                     pos_data['m'][i1],
                     pos_data['x'][i1],
                     pos_data['y'][i1],
                     pos_data['z'][i1],
                     pos_data['vx'][i1],
                     pos_data['vy'][i1],
                     pos_data['vz'][i1],
                     pos_data['eps'][i1],
                     pos_data['type'][i1],
                     pos_data['name'][i1],
                     pos_data['t'][i1],
                     pot, ax, ay, az, den
                     )
    out_file.write(out)
    
out_file.close()