import sys
from yt.mods import *
import numpy as np
from yt.analysis_modules.halo_finding.api import *
import glob
import re
import commands

#fn = 'output_%04i'%(output)
#fn = 'RedshiftOutput%04i'%output
dlist = glob.glob('RD00??')
dlist.sort()
#dlist.reverse() 
#n = int(re.sub(r'\D',"",fn))
for folder in dlist:
  n = int(re.sub(r'\D',"",folder))
  if glob.glob(folder+'/PopIII_haloes_%04i_type5.txt'%n) !=[]: continue
  fi1 = open(folder+'/sv.out','r')
  fi2 = open(folder+'/PopIII_list_%04i_type5.txt'%n,'r')
  fi3 = open(folder+'/PopIII_list_%04i_type1.txt'%n,'r') 
  fo1=open(folder+'/PopIII_haloes_%04i_type5.txt'%n,'w')
  fo2=open(folder+'/PopIII_haloes_%04i_type1.txt'%n,'w')

  line1 = fi1.readlines()
  line2 = fi2.readlines()
  line3 = fi3.readlines()
  lasthalo = len(line1)-1
  laststar2 = len(line2)-1
  laststar3 = len(line3)-1

  i=2
  while i<=lasthalo:
   hmass = float(line1[i].split()[1])
   hradius=float(line1[i].split()[13])
   hcx=float(line1[i].split()[7])
   hcy=float(line1[i].split()[8])
   hcz=float(line1[i].split()[9])
   fo1.write(str(i-2)+" "+str(hmass)+" "+str(hradius)+" "+str(hcx)+" "+str(hcy)+" "+str(hcz)+"\n")
   ct = [];mass=[];cx=[];cy=[];cz=[]
   j=1
   while j<=laststar2:
     sx=float(line2[j].split()[3])
     sy=float(line2[j].split()[4])
     sz=float(line2[j].split()[5]) 
     if ((hcx-sx)**2+(hcy-sy)**2+(hcz-sz)**2)**0.5<hradius:
         m=float(line2[j].split()[1])
         #if m > 1e20: m=m/1e20
         mass.append(m)
         ct.append(float(line2[j].split()[2]))
         cx.append(sx);cy.append(sy),cz.append(sz)
     j=j+1
   fo1.write("   "+str(len(mass))+"   "+str(np.mean(mass))+"\n")
   for j in xrange(len(mass)):
     fo1.write("   "+str(j)+" "+str(mass[j])+" "+str(ct[j])+" "+str(0)+" "+str(cx[j])+" "+str(cy[j])+" "+str(cz[j])+"\n")
   fo1.write("\n\n")
   fo2.write(str(i-2)+" "+str(hmass)+" "+str(hradius)+" "+str(hcx)+" "+str(hcy)+" "+str(hcz)+"\n")
   ct = [];mass=[];cx=[];cy=[];cz=[]
   j=1
   while j<=laststar3:
     sx=float(line3[j].split()[3])
     sy=float(line3[j].split()[4])
     sz=float(line3[j].split()[5])
     if ((hcx-sx)**2+(hcy-sy)**2+(hcz-sz)**2)**0.5<hradius:
         m=float(line3[j].split()[1])
         #if m > 1e20: m=m/1e20
         mass.append(m)
         ct.append(line3[j].split()[2])
         cx.append(sx);cy.append(sy),cz.append(sz)
     j=j+1
   fo2.write("   "+str(len(mass))+"   "+str(np.mean(mass))+"\n")
   for j in xrange(len(mass)):
     fo2.write("   "+str(j)+" "+str(mass[j])+" "+str(ct[j])+" "+str(0)+" "+str(cx[j])+" "+str(cy[j])+" "+str(cz[j])+"\n") 
   fo2.write("\n\n")
   i=i+1

  fi1.close()
  fi2.close()
  fi3.close()
  fo1.close()
  fo2.close()
