from yt.mods import *
from yt.analysis_modules.halo_profiler.api import *
from yt.analysis_modules.halo_finding.api import *
#from numpy import *
import numpy as np
import h5py

output = 15

pf = load("RedshiftOutput%04i"%output)

f = h5py.File('./Halos/RD%04i/VirialQuantities.h5'%output,'r')

data = f.get('RadiusMpc_200','f8')
Radius=np.asarray(data, dtype = 'f8')/pf.units['mpc']

data = f.get('center','f8')
center=np.asarray(data, dtype = 'f8')

data = f.get('id','i4')
id=np.asarray(data, dtype = 'i4')

f.close()

f=open('./halo_metallicity_%04i_virial.txt'%output,'a')

for i in xrange(Radius.size):
    sphere = pf.h.sphere(center[i], Radius[i]) 
    enriched_mass=sphere["CellMass"][sphere["TotalMetallicity"]>1e-4].sum()
    total_mass =sphere["CellMass"].sum()
    f.write(str(id[i])+" "+str(enriched_mass)+" "+str(total_mass)+"\n")   
   
f.close()