def pbload(filename, paramname=None):
    """
    Loads a snapshot using pynbody.  Can load a single species by appending
    ::gas, ::star, or ::dm to the filename
    
    Parameters
    ----------
    
    filename : str
        Filename to load
    paramname : str
        (optional) .param file to use
    
    Returns
    -------
    
    sim : snapshot
        A pynbody snapshot
    """
    if '::' in filename:
        
        filename, species = filename.split('::')
        sim = pb.load(filename, paramname=paramname)
        sim = getattr(sim, species)
        
    else:
        
        sim = pb.load(filename, paramname=paramname)
        
    return sim

def test_array_metadata() :
    f1 = pynbody.load("testdata/test_out.tipsy")
    
    f1.gas['zog'] = np.ones(len(f1.gas))
    f1.gas['zog'].units = "Msol kpc^-1"
    f1.gas['zog'].write()

    f1['banana'] = np.ones(len(f1))*0.5
    f1['banana'].units = "kpc^3 Myr^-1"
    f1['banana'].write()

    del f1

    f1 = pynbody.load("testdata/test_out.tipsy")
    assert "banana" in f1.loadable_keys()
    assert "zog" not in f1.loadable_keys()
    assert "banana" in f1.gas.loadable_keys()
    assert "zog" in f1.gas.loadable_keys()

    try:
        f1.star["zog"] # -> KeyError
        assert False # Shouldn't have been able to load gas-only array zog
    except KeyError :
        pass

    f1.gas['zog']
    assert f1.gas['zog'][0]==1.0
    assert f1.gas['zog'].units == "Msol kpc^-1"

    f1.star['banana']
    f1.gas['banana']
    f1.dm['banana']
    assert f1['banana'].units=="kpc^3 Myr^-1"

def test_unit_persistence():
    f = pynbody.load("testdata/test_g2_snap")

    # f2 is the comparison case - just load the whole
    # position array and convert it, simple
    f2 = pynbody.load("testdata/test_g2_snap")
    f2['pos']
    f2.physical_units()

    f.gas['pos']
    f.physical_units()
    assert (f.gas['pos'] == f2.gas['pos']).all()

    # the following lazy-loads should lead to the data being
    # auto-converted
    f.dm['pos']
    assert (f.gas['pos'] == f2.gas['pos']).all()
    assert (f.dm['pos'] == f2.dm['pos']).all()

    # the final one is the tricky one because this will trigger
    # an array promotion and hence internally inconsistent units
    f.star['pos']

    assert (f.star['pos'] == f2.star['pos']).all()

    # also check it hasn't messed up the other bits of the array!
    assert (f.gas['pos'] == f2.gas['pos']).all()
    assert (f.dm['pos'] == f2.dm['pos']).all()

    assert (f['pos'] == f2['pos']).all()

def test_snapshot_update() :
    f1 = pynbody.load("testdata/test_out.tipsy")
    f1['pos']
    f1['pos'] = np.arange(0,len(f1)*3).reshape(len(f1),3)
    
    # convert units -- the array should get written out in simulation units
    f1.g['pos'].convert_units('Mpc')

    f1['pos'].write(overwrite=True)
    f1.gas['metals'] = np.ones(len(f1.gas))*123.
    f1.star['metals'] = np.ones(len(f1.star))*345.

    f1.gas['metals'].write(overwrite=True)
    del f1

    f2 = pynbody.load("testdata/test_out.tipsy")
    assert (f2['pos']==np.arange(0,len(f2)*3).reshape(len(f2),3)).all()
    assert (f2.gas['metals']==123.).all() # should have updated gas.metals
    assert not (f2.star['metals']==345.).any() # should not have written out changes to star.metals

    # this is a completion:
    f2.dm['metals'] = np.ones(len(f2.dm))*789.1

    # should now be a simulation-level array... write it out
    f2['metals'].write(overwrite=True)

    del f2['metals']

    
    f3 = pynbody.load("testdata/test_out.tipsy")
    assert (f3.dm['metals']==789.1).all()

def test_fam_sim():
    """Check that an array loaded as families is the same as one loaded as a simulation array"""
    snap2 = pynbody.load("testdata/test_g2_snap")
    snap3 = pynbody.load("testdata/test_g2_snap")
    snap3.gas["pos"]
    snap3.dm["pos"]
    snap3.star["pos"]
    assert((snap3["pos"] == snap2["pos"]).all())

def test_units_override():
    f = pynbody.load("testdata/test_g2_snap.0")
    assert_equals(f.filename, "testdata/test_g2_snap")  # final '.0' is stripped
    assert_equals(f['pos'].units, "kpc a h^-1")

    # In this case the unit override system is not effective because
    # the final ".1" is not stripped away in the filename:
    # the file `test_g2_snap.units` is not used
    f_no_unit_override = pynbody.load("testdata/test_g2_snap.1")
    assert_equals(f_no_unit_override.filename, "testdata/test_g2_snap.1")
    assert_equals(f_no_unit_override['pos'].units, "Mpc a h^-1")  # from default_config.ini

def test_indexing():
    f1 = pynbody.load("testdata/g15784.lr.01024")

    test_len = 12310
    for test_len in [100, 10000, 20000]:
        for i in range(5):
            subindex = np.random.permutation(np.arange(0, len(f1)))[:test_len]
            subindex.sort()
            f2 = pynbody.load("testdata/g15784.lr.01024", take=subindex)

            assert (f2['x'] == f1[subindex]['x']).all()
            assert (f2['iord'] == f1[subindex]['iord']).all()

def getAllDust(lowz,catfiles,Rv,A1600max=2.0,halonum=[1],filename='dust.pkl'):
	f = open(catfiles,'r')
	files = f.readlines()
	dustExt = {'halos':halonum,'Rhl':np.zeros((len(halonum),len(files)+1)),'RDhalf':np.zeros((len(halonum),len(files)+1)),'z':np.zeros(len(files)+1),'u':np.zeros((len(halonum),len(files)+1)),'b':np.zeros((len(halonum),len(files)+1)),'v':np.zeros((len(halonum),len(files)+1)),'r':np.zeros((len(halonum),len(files)+1)),'i':np.zeros((len(halonum),len(files)+1)),'j':np.zeros((len(halonum),len(files)+1)),'h':np.zeros((len(halonum),len(files)+1)),'k':np.zeros((len(halonum),len(files)+1))}
	slz = pynbody.load(lowz)
        lz = str(round(slz.properties['a']**-1 -1,3))
        catend = '.cat.z'+lz+'\n'
        for i in range(len(files)):
                print "calculating dust corrections for ", files[i].strip('\n')
                xx = files[i].find(catend)
                simname=files[i][0:xx]
                s = pynbody.load(simname)
                h = s.halos()
                dustExt['z'][i] = s.properties['a']**-1 -1
                s.physical_units()
                catf = open(files[i].strip('\n'))
                cat = pickle.load(catf)
                catf.close()
		for j in range(len(halonum)):
                        print "halo", halonum[j]
                        progs, = np.where(cat==halonum[j])
                        if len(progs)==0:
                                print "no progenitors found in this step!"
                                continue
                        main = progs[0]
                        print "progenitor", main
                        h1 = h[main]
			dust,Rhalf = dustCor(h1,s,Rv,A1600max=A1600max)
			dustExt['RDhalf'][j,i] = Rhalf
#			dustExt['Rhl'][j,i] = Rhl
			for key in dust.keys():
				dustExt[key][j,i] = dust[key]
		del(s)
                del(h)
                del(h1)
                del(cat)
                gc.collect()
	print "calculating values for final step"
        h = slz.halos()
        slz.physical_units()
        for j in range(len(halonum)):
                h1 = h[halonum[j]]
		dust,Rhalf= dustCor(h1,slz,Rv)
		dustExt['z'][i+1] = slz.properties['a']**-1 -1
		dustExt['RDhalf'][j,i+1] = Rhalf
		for key in dust.keys():
			dustExt[key][j,i+1] = dust[key]
	if filename:
                print "saving data..."
                f = open(filename,'wb')
                pickle.dump(dustExt,f)
                f.close()
	return dustExt

def test_partial_loading() :
    f_f = pynbody.load("testdata/nchilada_test/12M.00001")
    
    test_ptcls = [ 11634,  24181,  26275,  37336,  37795,  38040,  38280,  38327,
        38524,  39349,  46758,  48892,  52160,  53267,  53745,  68970,
        78073,  83777,  86865,  93492,  94596,  96567,  99713, 106100,
       107856, 111036, 111830, 112560, 115082, 117111, 117444, 117667,
       123604, 123665, 124911, 132957, 138551, 154869, 158919, 182131,
       184252, 190498, 197946, 198288, 204526, 221720, 226375, 226915,
       229959, 231778] # randomly generated sample
    
    f_p = pynbody.load("testdata/nchilada_test/12M.00001", take = test_ptcls )

    assert((f_p['pos']==f_f['pos'][test_ptcls]).all())

def haloCat(lowz,highz,nhalos=50):
	print "matching halo catalog from ", highz, "to ", lowz
	s1 = pynbody.load(highz)
	s2 = pynbody.load(lowz)
	if s1.properties['a'] > s2.properties['a']:
		print "uh oh! highz file must actually be at higher z!"
		return
	b = pynbody.bridge.OrderBridge(s1,s2)
	cat = b.match_catalog()
	filename = highz+'.cat.z'+str(round(s2.properties['a']**-1-1,3))
	f = open(filename,'wb')
	pickle.dump(cat,f)
	f.close()
	return

def test_per_particle_loading():
    """Tests that loading one family at a time results in the
    same final array as loading all at once. There are a number of
    subtelties in the gadget handler that could mess this up by loading
    the wrong data."""

    f_all = pynbody.load("testdata/test_g2_snap")
    f_part = pynbody.load("testdata/test_g2_snap")

    f_part.dm['pos']
    f_part.star['pos']
    f_part.gas['pos']

    assert (f_all['pos'] == f_part['pos']).all()

def test_mass_unit_sanity():
    """Picks up on problems with converting array units as
    mass array gets loaded (which is a combination of a derived
    array and a loaded array)"""

    f1 = pynbody.load("testdata/ramses_partial_output_00250")
    f1['mass']
    f1.physical_units()

    f2 = pynbody.load("testdata/ramses_partial_output_00250")
    f2.physical_units()
    f2['mass']

    np.testing.assert_allclose(f1.dm['mass'], f2.dm['mass'], atol=1e-5)

def test_array_update():
    f1 = pynbody.load("testdata/test_out.tipsy")

    f1['bla'] = np.zeros(len(f1))
    f1['bla'].units = 'km'
    f1['bla'].write()

    del(f1['bla'])

    f1['bla']

    f1.g['bla'] = 1
    f1.d['bla'] = 2
    f1.s['bla'] = 3

    # test the case where bla is a snapshot-level array

    try:
        f1.g['bla'].write()
        assert False  # should not be allowed to overwrite here
    except IOError:
        pass

    f1.write_array(
        'bla', [pynbody.family.gas, pynbody.family.dm], overwrite=True)

    del(f1['bla'])

    f1['bla']

    assert all(f1.g['bla'] == 1)
    assert all(f1.d['bla'] == 2)
    assert all(f1.s['bla'] == 0)

    # test the case where bla2 is a family-level array

    f1.g['bla2'] = np.zeros(len(f1.g))
    f1.g['bla2'].units = 'km'
    f1.s['bla2'] = np.ones(len(f1.s))

    f1.write_array('bla2', [pynbody.family.gas, pynbody.family.star])

    del(f1)

    f1 = pynbody.load("testdata/test_out.tipsy")

    assert all(f1.g['bla2'] == 0)
    assert all(f1.s['bla2'] == 1)

def test_float_kd():
    f = pynbody.load("testdata/test_g2_snap")
    del f.properties['boxsize']

    assert f.dm['mass'].dtype==f.dm['pos'].dtype==np.float32
    assert f.dm['smooth'].dtype==np.float32

    # make double copy
    g = pynbody.new(len(f.dm))
    g.dm['pos']=f.dm['pos']
    g.dm['mass']=f.dm['mass']

    assert g.dm['mass'].dtype==g.dm['pos'].dtype==g.dm['smooth'].dtype==np.float64

    # check smoothing lengths agree (they have been calculated differently
    # using floating/double routines)

    npt.assert_allclose(f.dm['smooth'],g.dm['smooth'],rtol=1e-4)
    npt.assert_allclose(f.dm['rho'],g.dm['rho'],rtol=1e-4)

    # check all combinations of float/double smoothing
    double_ar = np.ones(len(f.dm),dtype=np.float64)
    float_ar = np.ones(len(f.dm),dtype=np.float32)

    double_double = g.dm.kdtree.sph_mean(double_ar,32)
    double_float = g.dm.kdtree.sph_mean(float_ar,32)
    float_double = f.dm.kdtree.sph_mean(double_ar,32)
    float_float = f.dm.kdtree.sph_mean(float_ar,32)

    # take double-double as 'gold standard' (though of course if any of these
    # fail it could also be a problem with the double-double case)

    npt.assert_allclose(double_double,double_float,rtol=1e-4)
    npt.assert_allclose(double_double,float_double,rtol=1e-4)
    npt.assert_allclose(double_double,float_float,rtol=1e-4)

def writeBHMark(simname,step,Name=None,iord=False,massrange=False):
	if not Name: f = open('BH.'+step+'.mark','w')
	else: f = open(Name,'w')
	s = pynbody.load(simname+'.'+step)
	f.write(str(len(s))+' '+str(len(s.gas))+' '+str(len(s.star))+'\n')
	if not iord: 
		if not massrange:
			bhind, = np.where(s.stars['tform']<0)
		else:
			if len(massrange) != 2:
				print "error massrange must be a length 2 tuple!"
				return
			bhind, = np.where((s.stars['tform']<0)&(s.stars['mass'].in_units('Msol')<massrange[1])&(s.stars['mass'].in_units('Msol')>massrange[0]))
	else:	
		bhind = np.array([])
		for ii in range(len(iord)):
			tmpind, = np.where(s.stars['iord']==iord[ii])
			if len(tmpind)==0: print "uh oh... iord ", iord[ii]," not found!"
			bhind = np.append(bhind,tmpind)
	bhindreal = bhind+len(s.dark)+len(s.gas)+1
	for ii in range(len(bhindreal)):
		f.write(str(bhindreal[ii])+'\n')
	f.close()
	del(s)
	return bhindreal

 def _LoadNewData(self, fluidtype):
     # Load data into pynbody
     print "STARTED LOADING DATA ", self._filename
     ro = self._snapdata
     if ro == None:
         try:
             ro=pynbody.load(self._filename)
         except:
             print "ABORTING LOADING (Opening snapshot failed - does this file exist?)"
             return
     if fluidtype == "stars":
         fluid = ro.stars
     elif fluidtype == "gas":
         fluid = ro.gas
     elif fluidtype == "dm":
         fluid = ro.dm
     else:
         print "Fluid type",fluidtype," not recognised! Use stars, gas or dm."
         raise TypeError
     posns = fluid["pos"]
     # No points?
     if len(posns) <= 1:
         print "ABORTING LOADING (No stars)"
         return
     smooth = fluid["smooth"]
     if not fluidtype == "gas":
         mass = fluid["mass"]
     else:
         mass = fluid["mass"]*fluid["temp"] # Hacky! Mmm. Basically display thermal energy in an element?
     if fluidtype == "dm":
         mass *= 0.1 # Artificially lower the brightness
     # Cut off background gas to save on rendering
     #if fluidtype == "gas":
     #    cutoff = 1e-3
     #    lim = mass > cutoff*np.max(mass)
     #    posns = posns[lim]
     #    smooth = smooth[lim]
     #    mass = mass[lim]
     # Process lengths to fit screen
     rescalepoints = True
     if rescalepoints:
         pmin, pmax = (np.min(posns),np.max(posns))
         posns = (posns - pmin) / (pmax - pmin)
         smooth /= (pmax - pmin)
     if smooth.max() <= 0.0:
         smooth *= 0.0
         smooth += 1.0
     else:
         smooth[smooth <= 0.0] = smooth[smooth > 0.0].min()
     cheaprescale = False
     if cheaprescale:
         cheap = fluid["pos"].units
         posns /= cheap
         smooth /= cheap
     posns -= 0.5
     self._outnum = 0#self._FindOutNum(filename)
     self._time = 0.0#ro._info["time"]
     self._snapdata = ro
     print "DONE LOADING DATA"
     return posns, smooth, mass

def load_center(output, align=True, halo=0):
    """
    Loads a RAMSES output and centers it on the desired halo. The hop
    center is used for an initial estimate, but for more precise
    centering, a shrinking-sphere center is calculated.

    **Inputs**:    
    
    *output* : path to RAMSES output directory

    **Optional Keywords**: 

    *align* : whether to align the snapshot based on the angular momentum in the central region (default = True)

    *halo* : halo to center on (default = 0)
    """


    s = pynbody.load(output)
    hop_center(s,halo)

    st = s[pynbody.filt.Sphere('100 kpc')]
    
    cen = pynbody.analysis.halo.center(st,retcen=True,mode='ssc',verbose=True)
    
    if align: 
        pynbody.analysis.angmom.faceon(st.s,disk_size='10 kpc',cen=cen,mode='ssc')
    else :
        s['pos'] -= cen

    s['pos'].convert_units('kpc')
    s['vel'].convert_units('km s^-1')

    return s

def test_alt_names():
    _h5py_copy_with_key_rename('testdata/Test_NOSN_NOZCOOL_L010N0128/data/snapshot_103/snap_103.hdf5',
                'testdata/Test_NOSN_NOZCOOL_L010N0128/data/snapshot_103/snap_103_altnames.hdf5')

    snap_alt = pynbody.load('testdata/Test_NOSN_NOZCOOL_L010N0128/data/snapshot_103/snap_103_altnames.hdf5')
    assert 'mass' in snap_alt.loadable_keys()
    assert all(snap_alt['mass']==snap['mass'])

def read_full_orbit_file(filename, simname):
	f = open('files.list','r')
	files = f.readlines()
	s = pynbody.load(files[0].strip('\n'))
	f.close()
	a = None

	try:
		bhid, time, step, mass, x, y, z, vx, vy, vz, pot, mdot, dm, E, dt, a = readcol(filename, twod=False)
	except:
		bhid, time, step, mass, x, y, z, vx, vy, vz, pot, mdot, dm, E, dt = readcol(filename, twod=False)

	output = {'iord':bhid, 'time':time, 'step':step, 'mass':mass, 'x':x, 'y':y,
			  'z':z, 'vx':vx, 'vy':vy, 'vz':vz, 'pot':pot, 'mdot':mdot, 'dm':dm, 'E':E, 'dt':dt, 'a':a}
	if a is None:
		a, = cosmology.getScaleFactor(pynbody.array.SimArray(time,s.infer_original_units('Gyr')),s)
		output['a'] = a
	units = {'x':'kpc', 'y':'kpc', 'z':'kpc', 'vx':'km s**-1', 'vy':'km s**-1', 'vz':'km s**-1',
			 'mdot':'Msol yr**-1', 'dm':'Msol', 'dt':'Gyr', 'time':'Gyr', 'mass':'Msol'}

	for key in ['x', 'y', 'z', 'vx', 'vy', 'vz']:
		print "fixing scale factor for", key
		output[key] *= output['a']

	for key in units.keys():
		print "converting units for ", key
		origunit = s.infer_original_units(units[key])
		if key in ['x', 'y', 'z', 'vx', 'vy', 'vz']: origunit = origunit / pynbody.units.Unit('a')
		output[key] = pynbody.array.SimArray(output[key],origunit)
		output[key] = output[key].in_units(units[key])

	order = np.argsort(output['time'])
	util.cutdict(output,order)
	return output

	def gethalos(self):
		f = open('files.list', 'r')
		simfiles = f.readlines()
		nsnaps = len(simfiles)
		f.close()
		f = open('steps', 'r')
		snaps = f.readlines()
		f.close()
		initarr = np.ones((len(self.bhiords), len(snaps))) * -1
		self.bhhalos = {'Grp': initarr}
		for i in range(nsnaps):
			if os.path.exists(simfiles[i].strip('\n') + '.rockstar.grp'):
				grpfile = simfiles[i].strip('\n') + '.rockstar.grp'
			else:
				if os.path.exists(simfiles[i].strip('\n') + '.amiga.grp'):
					grpfile = simfiles[i].strip('\n') + '.amiga.grp'
				else:
					print "ERROR there is no grp file for this step!"
					continue
			sim = pynbody.load(simfiles[i].strip('\n'))
			simind = np.where(np.in1d(sim.stars['iord'], self.bhiords))
			orbind = np.where(np.inqd(self.bhiords, sim.stars['iord']))
			simind += len(sim.dm) + len(sim.gas)
			del sim['iord']
			gc.collect()
			grp = readcol(grpfile, skipline=1)
			self.bhhalos['Grp'][orbind][:, i] = grp[simind]
			del grp, simind, orbind
			gc.collect()