from pylab import *

import os

import scipy.constants as consts

G = consts.value("Newtonian constant of gravitation")
c = consts.value("speed of light in vacuum")
thomson = consts.value("Thomson cross section")
mproton = consts.value("proton mass")
melectron = consts.value("electron mass")
mhenucleus = consts.value("alpha particle mass")
m_h = mproton + melectron        # neglecting binding energy
m_heII = mhenucleus + melectron  # neglecting binding energy
kb = consts.value("Boltzmann constant")
pc = consts.parsec

Gpc = 1.0e9 * pc
Mpc = 1.0e6 * pc
kpc = 1.0e3 * pc

km = 1000.0
cm_SI = 0.01
angstrom = 1.0e-10
s = 1.0

msun = 1.9884430e30 # in kg

xH = 0.76


class spectra:
  def __init__(self, spectype, filename, nspec=0, taufilename="none"):
    
    if spectype in ["se","se_onthefly"]:
      specfile = open(filename, 'rb')
      
      # header
      self.z = fromfile(specfile, dtype=float64, count=1)[0]
      self.om = fromfile(specfile, dtype=float64, count=1)[0]
      self.ol = fromfile(specfile, dtype=float64, count=1)[0]
      self.ob = fromfile(specfile, dtype=float64, count=1)[0]
      self.h = fromfile(specfile, dtype=float64, count=1)[0]
      self.box = fromfile(specfile, dtype=float64, count=1)[0]
      self.xh = fromfile(specfile, dtype=float64, count=1)[0]     # hydrogen mass fraxtion
      self.nbins = fromfile(specfile, dtype=int32, count=1)[0]
      self.nlos = fromfile(specfile, dtype=int32, count=1)[0]
      
      self.a = 1.0/(1.0+self.z)
      
      self.H0 = self.h * 100.0 * km / s / Mpc   # SI
      
      # postions
      self.dirlos = fromfile(specfile, dtype=int32, count=self.nlos)
      self.xlos = fromfile(specfile, dtype=float64, count=self.nlos)
      self.ylos = fromfile(specfile, dtype=float64, count=self.nlos)
      self.zlos = fromfile(specfile, dtype=float64, count=self.nlos)
      
      # coordinate and velocity of pixel center
      self.pixpos = fromfile(specfile, dtype=float64, count=self.nbins) # first pixel center is at simulation coordinate 0 and has pixpos = 0
      self.pixvel = fromfile(specfile, dtype=float64, count=self.nbins) # pixpos * a * H(a)
      
      self.dxpix = self.box/self.nbins
      self.dxpix_phys = self.dxpix*self.a
      
      assert self.box > 1000.0 # otherwise likely not kpc/h units
      self.dxpix_phys_SI =self.dxpix_phys * kpc / self.h # assumes kpc / h units
      
      self.vmax = (self.pixvel[-1]-self.pixvel[0])/(self.nbins-1)*self.nbins
      self.dvpix = self.pixvel[1]-self.pixvel[0]
      
      # gas overdensity
      if spectype=="se":
        #self.rhoH2rhoHmean = 10.0**fromfile(specfile, dtype=float64, count=self.nbins*self.nlos).reshape(self.nlos,self.nbins)
        self.rhoH2rhoHmean = fromfile(specfile, dtype=float64, count=self.nbins*self.nlos).reshape(self.nlos,self.nbins)
      elif spectype=="se_onthefly":
        self.rhoH2rhoHmean = fromfile(specfile, dtype=float64, count=self.nbins*self.nlos).reshape(self.nlos,self.nbins)
      
      # HI Lyman-alpha
      self.nHI_frac = fromfile(specfile, dtype=float64, count=self.nbins*self.nlos).reshape(self.nlos,self.nbins)
      self.temp_HI = fromfile(specfile, dtype=float64, count=self.nbins*self.nlos).reshape(self.nlos,self.nbins)
      self.vel_HI = fromfile(specfile, dtype=float64, count=self.nbins*self.nlos).reshape(self.nlos,self.nbins)
      
      if spectype=="se":
        self.tau_HI = fromfile(specfile, dtype=float64, count=self.nbins*self.nlos).reshape(self.nlos,self.nbins)
      elif spectype=="se_onthefly":
        if taufilename!="none":
          self.tau_HI = fromfile(taufilename, dtype=float64, count=self.nbins*self.nlos).reshape(self.nlos,self.nbins)   
      
      curpos = specfile.tell()
      specfile.seek(0,os.SEEK_END)
      endpos = specfile.tell()
      
      self.have_HeII = False
      
      if curpos != endpos:
        specfile.seek(curpos,os.SEEK_SET)
        
        ## HeII Lyman-alpha
        self.nHeII_frac = fromfile(specfile, dtype=float64, count=self.nbins*self.nlos).reshape(self.nlos,self.nbins)
        self.temp_HeII = fromfile(specfile, dtype=float64, count=self.nbins*self.nlos).reshape(self.nlos,self.nbins)
        self.vel_HeII = fromfile(specfile, dtype=float64, count=self.nbins*self.nlos).reshape(self.nlos,self.nbins)
        self.tau_HeII = fromfile(specfile, dtype=float64, count=self.nbins*self.nlos).reshape(self.nlos,self.nbins)
      
        self.have_HeII = True
      
      curpos = specfile.tell()
      if curpos != endpos:
        print "WARNING: not reading the whole file", curpos, "of", endpos
      
      specfile.close()
    
    else:
      assert False
      
      

spec_obj = spectra("se_onthefly","los2048_n5000_z0.000.dat")
Delta_arr = spec_obj.rhoH2rhoHmean.T
T_gas_arr = spec_obj.temp_HI.T
v_gas_arr = spec_obj.vel_HI.T
f_HI_arr  = spec_obj.nHI_frac.T # nHI/nH
print Delta_arr.shape
print T_gas_arr.shape
print v_gas_arr.shape
print f_HI_arr.shape