#--------------------------------------------importing packages--------------------------------------------
import bigfile
import numpy
import os
import shutil
#--------------------------------------------create directory for storing the snapshot if does not exist--- 
def ensure_dir(f):
    d = os.path.dirname(f)
    if not os.path.exists(d):
        os.makedirs(d)
    else:
        shutil.rmtree(f)
        os.makedirs(d)
z=7.5        
f='/nfs/nas-0-1/akbhowmi/BT2_rockstar_galaxies/snapshots/PIG_mulFOF1000_141/'
ensure_dir(f)
#----------------------------------------------------------------------------------------------------------
#------------------------------------------------open bigfile objects for source and destination-----------
pig=bigfile.BigFile('/physics2/vat/bluetide2_hydro/PIG_141/')
pig_new=bigfile.BigFile('/home/akbhowmi/PIG_141')


pig2=bigfile.BigFile(f,create=True)
#----------------------------------------------------------------------------------------------------------

#-----------------------------------------Read FOF id------------------------------------------------------



#fofid_i=0 #---------------initial fof group----------------------------------------------------------------
N=1000 #---------------------no. of FOF groups----------------------------------------------------------------

f=open('/nfs/nas-0-1/akbhowmi/rockstar_run_BT2/job_folder/read_index.txt','r')
fofid_i=int(f.read())*1000
f.close()



#fofid_i=0

#----------------------------------------------------------------------------------------------------------
#----------------------------------------Generate the bigfile blocks---------------------------------------
for f in pig.list_blocks():
    if (('0' in f) or ('1' in f)) or (('4' in f) or ('5' in f)): 
        FOFlength=pig.open('/FOFGroups/LengthByType')
        FOFOffset=pig_new.open('/OffsetByType')
	if(((('Position' in f) or ('Velocity' in f)) or (('ID' in f) or ('InternalEnergy' in f))) or ('Mass' in f)):
        	if('0' in f):
            		FOFlength=FOFlength[fofid_i+N][0]
            		FOFOffset_i=FOFOffset[fofid_i][0]
	    		FOFOffset_f=FOFOffset[fofid_i+N][0]    
        	elif('1' in f):
            		FOFlength=FOFlength[fofid_i+N][1]
            		FOFOffset_i=FOFOffset[fofid_i][1]
	    		FOFOffset_f=FOFOffset[fofid_i+N][1]
        	elif('4' in f):
            		FOFlength=FOFlength[fofid_i+N][4]
            		FOFOffset_i=FOFOffset[fofid_i][4]
	    		FOFOffset_f=FOFOffset[fofid_i+N][4]
        	elif('5'in f):
            		FOFlength=FOFlength[fofid_i+N][5]
            		FOFOffset_i=FOFOffset[fofid_i][5]
	    		FOFOffset_f=FOFOffset[fofid_i+N][5]
		length=FOFOffset_f+FOFlength-FOFOffset_i
            
        	pig2.create(f,dtype=(pig.open(f)).dtype,size=length,Nfile=8)
		print 'reading ', f
		#arr=bigfile.BigBlock.read(pig.open(f),FOFOffset_i,length)
        	arr=pig.open(f)[FOFOffset_i:FOFOffset_i+length]
		print 'writing ', f
		bigfile.BigBlock.write(pig2.open(f),0,arr)


#------------------------------------------Creating header file---------------------------------------------------

hi=pig.open('Header')

h=pig2.create('header')

p=[pig2.open('%d/Mass'%i) for i in [0,1,4,5]]

for i in hi.attrs.keys():
    att_val=hi.attrs.__getitem__(i)
    if ('NumPart' not in i ):
        h.attrs.__setitem__(i,att_val)
    if ('NumPart' in i):
        h.attrs.__setitem__(i,numpy.array([p[0].size,p[1].size,0,0,p[2].size,p[3].size]))
h.close()
pig2.close()
pig.close()