from math import log10
import math
import numpy as na
#nullfmt = NullFormatter()
#from mpi4py import MPI
import re
import glob
import os
import time

#comm = MPI.COMM_WORLD
#rank = comm.Get_rank()
#size=comm.size

redshift = 15.0
#redshift_end = 23.5
#tfinal = 4345107   #4615649   # in years
nx = 512
Ex = 300.0   # photon energy (eV)
cfl = 0.5
bmin=10.0
bmax=1500.0
bins = 100

lengthunit = 40.0*1000.0/(redshift+1.0)/512.0
bmin = bmin / lengthunit
bmax = bmax / lengthunit
bwidth = math.log10(bmax/bmin)/bins

def array_2_profile(arr,nx,bins,bmin,bmax,bwidth):
    profile = na.zeros(bins)
    weight =  na.zeros(bins)
    for ii in xrange(nx):
       for jj in xrange(nx):
           for kk in xrange(nx):
              index1 = int(math.log10(max(((ii-nx/2.0)**2+(jj-nx/2.0)**2+(kk-nx/2.0)**2)**0.5/bmin,0.01))/bwidth+0.5)
              if index1 >= bins: continue
              profile[index1] += arr[ii][jj][kk]
              weight[index1] += 1
    return profile/weight


filelist = glob.glob("Final_Electron_fraction_*%5.2f.dat"%redshift)
filelist.sort()
#xraybins = len(filelist)
#xrayzstart = na.zeros([xraybins])
#xrayzend = na.zeros([xraybins])
#xrayenergy = na.zeros([xraybins])
tem_profile = []
#Xray_r = []
for i,name in enumerate(filelist):
  #xrayzstart[i] = float(re.findall(r"[-+]?\d*\.\d+|\d+",name)[0])
  #xrayzend[i] = float(re.findall(r"[-+]?\d*\.\d+|\d+",name)[1])
  #xrayenergy[i] = Ex*(1+redshift)/(1+(xrayzstart[i]+xrayzend[i])/2)
  f = open(name,'r')
  #Xray_r.append(na.fromfile(f,dtype='float64').reshape(nx,nx,nx))
  Tem=na.fromfile(f,dtype='float64').reshape(nx,nx,nx)
  print "Reading %s"%name 
  f.close()
  tem_profile.append(array_2_profile(Tem,nx,bins,bmin,bmax,bwidth))


#del rho,Xray_r
#del Tem, ef

fo=open('ef-profile_%5.2f_new.txt'%redshift,'w')
fo.write('radius (kpc)   Electron fraction')
fo.write('\n')
for i in xrange(bins):
   fo.write(str(10**((i-0.5)*bwidth)*bmin*lengthunit)+" ")
   fo.write(str(tem_profile[0][i])+"  ")
   fo.write("\n")