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

tunit = 6.88027e+14
tyr = tunit/3.15e7
omega_b = 0.0449
h = 0.71
baryonmass = 1.9e-29*h**2*omega_b*((8*3.08e24)**3*0.2745822715872526)/2e33
mass_solar = 1.998e33
L_ed = 1.3e38 * 6.24e11 # ev
epsilon = 0.2
T_ed = epsilon*mass_solar/1.44e17
e_ph0 = 4.202e3  # 1.2 ev Temperature, median wavelength
skip = 0 # half PopIII stars become binaries
lifetime_max = 30.0 #Myr
mass_min = 40
mass_min1 = 10

dlist = glob.glob('RD00??')
dlist.sort()
#output = int(re.sub(r'\D',"",fn))
for folder in dlist:
  lifetime_max = 30.0
  n = int(re.sub(r'\D',"",folder))
  if glob.glob(folder+'/PopIII_haloes_xray-sources_distribution_%02iMyr_40Ms.txt'%lifetime_max) !=[]: continue
  a=commands.getoutput('grep InitialTime '+folder+'/RedshiftOutput00??')
  currenttime=float(a.split()[2])
  f=open(folder+'/PopIII_haloes_%04i_type1.txt'%n,'r')
  f1=open(folder+'/PopIII_haloes_%04i_type5.txt'%n,'r')
  f2=open(folder+'/Overlaped_haloes_list.txt','r')
  overlayed=[]
  line=f2.readline()
  while line !='':
   overlayed.append(int(line))
   line=f2.readline()
  f2.close()
  fo = open(folder+'/PopIII_haloes_xray-sources_distribution_%02iMyr_40Ms.txt'%lifetime_max,'w')

  lifetime_max = lifetime_max*1e6
  mass_max = mass_min*np.exp(lifetime_max*3.15e7/T_ed)
  mass_max1 = mass_min1*np.exp(lifetime_max*3.15e7/T_ed)
  emission_max = epsilon*(mass_max-mass_min)*mass_solar*(3.0e10)**2*6.24e11
  emission_max1 = epsilon*(mass_max1-mass_min1)*mass_solar*(3.0e10)**2*6.24e11
  source_number = 0

  line = f.readline()
  line1 = f1.readline()
  j=0
  while line !='':
   halomass = float(line.split()[1])
   haloposition = float(line.split()[3]), float(line.split()[4]),float(line.split()[5].replace(":",""))
   energyphoton = []; lifetime = []; emission = []; creationtime = []
   line = f.readline();line = f.readline()
   while line !="\n":
     skip = skip + 1
     starmass = float(line.split()[1])
     if skip%2 == 0:
      if starmass > mass_min:
        creationtime.append(float(line.split()[2]))
        energyphoton.append(e_ph0*(10.0/(starmass+mass_min)*2)**0.25)
        lt = T_ed*np.log((starmass)/mass_min)
        lifetime.append(lt if lt < lifetime_max*3.15e7 else lifetime_max*3.15e7)
        em = epsilon*(starmass-mass_min)*mass_solar*(3.0e10)**2*6.24e11
        emission.append(em if starmass < mass_max else emission_max) 
      elif starmass > mass_min1:
        creationtime.append(float(line.split()[2]))
        energyphoton.append(e_ph0*(10.0/(starmass+mass_min1)*2)**0.25)
        lt = T_ed*np.log((starmass)/mass_min1)
        lifetime.append(lt if lt < lifetime_max*3.15e7 else lifetime_max*3.15e7)
        em = epsilon*(starmass-mass_min1)*mass_solar*(3.0e10)**2*6.24e11
        emission.append(em if starmass < mass_max1 else emission_max1)    
     line=f.readline()
   line1 = f1.readline();line1 = f1.readline()
   while line1 !="\n":
     skip = skip + 1
     starmass = float(line1.split()[1])
     if skip%2 == 0:
      if starmass > 1e20: starmass = starmass/1.0e20
      if starmass > mass_min:
        creationtime.append(float(line1.split()[2]))
        energyphoton.append(e_ph0*(10.0/(starmass+mass_min)*2)**0.25)
        lt = T_ed*np.log((starmass)/mass_min)
        lifetime.append(lt if lt < lifetime_max*3.15e7 else lifetime_max*3.15e7)
        em = epsilon*(starmass-mass_min)*mass_solar*(3.0e10)**2*6.24e11
        emission.append(em if starmass < mass_max else emission_max) 
      elif starmass > mass_min1:
        creationtime.append(float(line1.split()[2]))
        energyphoton.append(e_ph0*(10.0/(starmass+mass_min1)*2)**0.25)
        lt = T_ed*np.log((starmass)/mass_min1)
        lifetime.append(lt if lt < lifetime_max*3.15e7 else lifetime_max*3.15e7)
        em = epsilon*(starmass-mass_min1)*mass_solar*(3.0e10)**2*6.24e11
        emission.append(em if starmass < mass_max1 else emission_max1)  
     line1=f1.readline() 
   energyphoton = np.array(energyphoton); lifetime = np.array(lifetime); 
   emission = np.array(emission) ; creationtime = np.array(creationtime)
   if energyphoton.size != 0:
     emission = np.array(emission)
     #emission[emission>=emission_max]=emission_max
     totalemission = emission.sum()
     lifetime = np.array(lifetime)
     #lifetime[lifetime >= lifetime_max*3.15e7] =lifetime_max*3.15e7 
     halolifetime = lifetime.mean()
     luminosity = emission/lifetime
     luminosity[(currenttime-creationtime)>(lifetime/tunit)]=0
     halolifetime = halolifetime/tunit     
     halocreationtime = creationtime.min()
     luminosity_sum=luminosity.sum()
     if j not in overlayed:
        fo.write(str(halomass)+" "+str(luminosity_sum)+" "+str((luminosity>0).sum())+"\n")   
        source_number = source_number + 1
   j = j +1
   line = f.readline()
   if line == "\n": line = f.readline()
   line1 = f1.readline();
   if line1 == "\n": line1 = f1.readline()

  f.close()
  f1.close()
  fo.close()