import yt
import glob
import sys
import os
from yt.utilities.physical_ratios import rho_crit_g_cm3_h2
import yt.units as u
import numpy as np
import mag_initialize as m
import mag_write_data_functions as mwdf

h = 0.7
Omega_m = 0.27 
Omega_Lambda = 0.73
rho_c = rho_crit_g_cm3_h2*h*h*u.g/u.cm**3
rho_c.convert_to_units("Msun/kpc**3")
                                                                          
def find_center(ds, halo=1):
    print('Finding gravitational potential minimum')
    # Find the gravitational potential minimum                                                                                                                                                                                                             
    #####                                                                                                                                                                                                                                                  
    # Get all the data                                                                                                                                                                                                                                     
    dd = ds.all_data()

    # Find the halo's particles                                                                                                                                                                                                         
    logic = dd["particle_halo"].astype("int") == halo

    idx = np.argsort(dd['io','particle_gpot']*logic)[:100]
    c = dd["particle_position"][idx,:].mean(axis=0)

    return c

small = 1.0e-5

def find_mr500(ds, c):
    def _flash_mass(field, data):
        return (data['dens']+data['pden'])*data['cell_volume']
    ds.add_field(("gas","flash_mass"), _flash_mass, units="g")
    rhi = ds.quan(2500.0, "kpc")
    rlo = ds.quan(0.0, "kpc")
    sp = ds.sphere(c, rhi)
    r = sp["radius"].to('kpc').d
    idxs = np.argsort(r)
    m = np.cumsum(sp["flash_mass"][idxs].to('Msun').d)
    r = r[idxs]
    rhoavg = (3.*m/(4.*np.pi*r**3))
    idx = np.searchsorted(rhoavg[::-1], 500.0*rho_c.v) - 1
    return m[::-1][idx], r[::-1][idx]

halo = int(sys.argv[2])
root = "/data/mimir/jzuhone/data/fid/"
folder = sys.argv[1]
basenm = "fiducial_%s_hdf5_plt_cnt_0[0-9]*0" % folder

files = os.path.join(root, folder, basenm)

print(files)

fns = glob.glob(files)
fns.sort()

ts = yt.DatasetSeries(fns)

fid = []
times = []
x = []
y = []
z = []
m500c = []
r500c = []
for ds in ts:
    print(ds)
    c = find_center(ds, halo=halo)
    print(c)
    m, r = find_mr500(ds, c)
    fid.append(int(str(ds)[-4:]))
    x.append(c[0])
    y.append(c[1])
    z.append(c[2])
    times.append(ds.current_time.to("Gyr").v)
    m500c.append(m)
    r500c.append(r)

np.savetxt("%s_halo%d_mr500.dat" % (folder, halo),
           np.transpose([fid, times, x, y, z, r500c, m500c]),
           delimiter="\t", header=" id\tt\tx\ty\tz\tr\tM", 
           fmt='%d\t%.18e\t%.18e\t%.18e\t%.18e\t%.18e\t%.18e')