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
mass_min = 40.0
mass_min1 = 10.0
lifetime_max = 3.0e7 # 30 Myr
mass_max = mass_min*np.exp(lifetime_max*3.15e7/T_ed)  # accrete at Eddington limit for max lifetime
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
e_ph = e_ph0*(10.0/(40.0))**0.25   #e_ph0*(10.0/(mass_max+mass_min)*2)**0.25


dlist = glob.glob('RD00??')
dlist.sort()
for folder in dlist:
  n = int(re.sub(r'\D',"",folder))
  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_lifetime_40Ms.txt','w')

  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
     if skip%2 == 0:
      starmass = float(line.split()[1])
      if starmass > mass_min:
        ct = float(line.split()[2])
        lt = T_ed*np.log((starmass)/mass_min)
        if lt > lifetime_max*3.15e7: lt = lifetime_max*3.15e7
        if (currenttime - ct) < lt/tunit:
          creationtime.append(float(ct))
          energyphoton.append(e_ph)
          lifetime.append(lt)        
          emission.append(epsilon*((starmass if starmass < mass_max else mass_max)-mass_min)*mass_solar*(3.0e10)**2*6.24e11)  
      elif starmass > mass_min1:
        ct = float(line.split()[2])
        lt = T_ed*np.log((starmass)/mass_min1)
        if lt > lifetime_max*3.15e7: lt = lifetime_max*3.15e7
        if (currenttime - ct) < lt/tunit:
          creationtime.append(float(ct))
          energyphoton.append(e_ph)
          lifetime.append(lt)
          emission.append(epsilon*((starmass if starmass < mass_max1 else mass_max1)-mass_min1)*mass_solar*(3.0e10)**2*6.24e11) 
     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:
       ct = float(line1.split()[2])
       lt = T_ed*np.log((starmass)/mass_min)
       if lt > lifetime_max*3.15e7: lt = lifetime_max*3.15e7
       if (currenttime - ct) < lt/tunit:
          creationtime.append(float(ct))
          energyphoton.append(e_ph)
          lifetime.append(lt)
          emission.append(epsilon*((starmass if starmass < mass_max else mass_max)-mass_min)*mass_solar*(3.0e10)**2*6.24e11)
      elif starmass > mass_min1:
       ct = float(line1.split()[2])
       lt = T_ed*np.log((starmass)/mass_min1)
       if lt > lifetime_max*3.15e7: lt = lifetime_max*3.15e7
       if (currenttime - ct) < lt/tunit:
          creationtime.append(float(ct))
          energyphoton.append(e_ph)
          lifetime.append(lt)
          emission.append(epsilon*((starmass if starmass < mass_max1 else mass_max1)-mass_min1)*mass_solar*(3.0e10)**2*6.24e11)
     line1=f1.readline() 
   energyphoton = np.array(energyphoton); lifetime = np.array(lifetime); 
   emission = np.array(emission) ; creationtime = np.array(creationtime)
   if np.logical_and(energyphoton.size != 0 , j not in overlayed):
     totalemission = emission.sum()
     halolifetime = lifetime.max()
     luminosity = totalemission/halolifetime
     halophotonenergy = (energyphoton*emission).sum()/emission.sum()
     luminosity = luminosity/halophotonenergy
     halolifetime = halolifetime/tunit     
     halocreationtime = creationtime.mean()
     fo.write(str("PhotonTestSourceType[%d]       = 1"%source_number+"\n"))
     fo.write(str("PhotonTestSourcePosition[%d]   = %f %f %f"%(source_number, \
        haloposition[0],haloposition[1],haloposition[2])+"\n"))
     fo.write(str("PhotonTestSourceLuminosity[%d] = %e"%(source_number,luminosity)+"\n"))
     fo.write(str("PhotonTestSourceLifeTime[%d]   = %f"%(source_number,halolifetime)+"\n"))
     fo.write(str("PhotonTestSourceCreationTime[%d] = %f"%(source_number,halocreationtime)+"\n"))
     fo.write(str("PhotonTestSourceEnergyBins[%d] = 1"%source_number+"\n")) 
     fo.write(str("PhotonTestSourceEnergy[%d] = %f"%(source_number,halophotonenergy)+"\n\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()