import os
import numpy as np
import re
import sys
import fitsio
import maskutils

def getnbarbychunk(catfname,maskfname,randomfname=None,nbarfname=None):
  """
  Returns ratio of data to randoms chunk-by-chunk, relative to the ratio in the full sample.
  This should be an unbiased estimate of nbar (relative to the full sample).
  However, this requires randoms.  In the case where we didn't produce randoms yet, let's just 
  take raw galaxies and divide by mask area.  This will be systematically low because of veto'ing, but gives a rough guess. 

  Usage case 1: catfname = target catalog, maskfname = binarymask
  Usage case 2: catfname = data catalog, randomfname = random catalog, maskfname = completeness mask
  """

  nbarchunk = np.zeros(39)
  Nchunk = np.zeros(39)
  nbarchunksys = np.zeros(39)
  dd = fitsio.read(catfname)
  mm = maskutils.mask(maskfname)

  ## trim the galaxy catalog by the binary mask; probably not necessary.
  xx = np.where(mm.weight[dd['ISECT']] > 0.001)[0]
  dd = dd[xx]

  myavg = -1. # don't compute for rough estimate with target cat.

  if randomfname is not None:
    assert nbarfname is not None ## need this to get total effective area in mask, for conversion of # of randoms t o area.
    ifpn = open(nbarfname,'r')
    line = ifpn.readline()
    ll = line.strip('\n').split(': ')[1].split(' ')
    for i in range(ll.count('')):
      ll.remove('')
    aefftot = float(ll[0])
    vefftot = float(ll[1])
    myavg = (dd['WEIGHT_CP'][:] + dd['WEIGHT_NOZ'][:] - 1.0).sum()/aefftot
    rr = fitsio.read(randomfname)
    xx = np.where(mm.weight[rr['ISECT']] > 0.001)[0]
    rr = rr[xx]
    nrantot = len(rr)
    rtoa = aefftot/nrantot
    meansys = (dd['WEIGHT_SYSTOT'][:]*(dd['WEIGHT_CP'][:] + dd['WEIGHT_NOZ'][:] - 1.0)).sum()/((dd['WEIGHT_CP'][:] + dd['WEIGHT_NOZ'][:] - 1.0)).sum()
    print 'meansys',meansys
  for chunk in range(2,39):
    xx = np.where(mm.mf['CHUNKVECTOR'][:,0] == chunk)[0]
    aeff = (mm.mf['STR'][xx]*mm.mf['WEIGHT'][xx]).sum()*(180./np.pi)**2
    goodsec = mm.mf['SECTOR'][xx]
    isec = np.zeros(len(mm.area),dtype='int')
    isec[goodsec] = 1 
    aeffchk = (mm.area[:]*isec[:]*mm.weight[:]).sum()*(180./np.pi)**2
    if aeff > 0:
      assert np.fabs(aeff - aeffchk)/aeff < 2.0e-6
    if randomfname == None:
      xx = np.where(isec[dd['ISECT']] == 1)[0]
      nbarchunk[chunk] = (len(xx))/aeff
      Nchunk[chunk] = len(xx)
      print chunk, len(xx), nbarchunk[chunk] 
    else:
      ## get number of galaxies in the sector.
      xx = np.where(isec[dd['ISECT']] == 1)[0]
      xxr = np.where(isec[rr['ISECT']] == 1)[0]
      ## get area weighted completeness.
      aeffnew = len(xxr)*rtoa
      Nchunk[chunk] = (dd['WEIGHT_CP'][xx] + dd['WEIGHT_NOZ'][xx] - 1.0).sum()
      nbarchunk[chunk] = Nchunk[chunk]/aeffnew
      nbarchunksys[chunk] = (dd['WEIGHT_SYSTOT'][xx]*(dd['WEIGHT_CP'][xx] + dd['WEIGHT_NOZ'][xx] - 1.0)).sum()/aeffnew/meansys
      print chunk, len(xx), nbarchunk[chunk], nbarchunksys[chunk]

  return Nchunk, nbarchunk, nbarchunksys, myavg
    ##


def getdr12v4nbar(sampletag='cmass',NorS=0):
  """
  Hard-coding the veto fractions and file paths for each sample.  Ick!!
  Do this better later.
  Returns number density in units of galaxy targets per steradians.
  """

  assert sampletag == 'cmass' or sampletag == 'lowz'

  if NorS == 0:
    NS = 'N'
  if NorS == 1:
    NS = 'S'

  if sampletag == 'cmass':
    mb = maskutils.mask("/home/bareid/boss/bosslss/trunk/data/binarymask-dr12v4-"+NS+".fits")

  else:
    mb = maskutils.mask("/home/bareid/boss/bosslss/trunk/data/binarymask-dr12v4-lowzchunkcut-"+NS+".fits")

  targfname = "/home/bareid/boss/bosslss/trunk/data/"+sampletag+"-dr12v4-"+NS+"-Reid-targ.dat.fits"
  dd = fitsio.read(targfname)
  wgttot = dd['WEIGHT_SYSTOT'][:].sum()
  print 'hi beth, avg syswgt:',wgttot/float(len(dd))

  if sampletag == 'cmass' and NorS == 0:
    vetof = 0.5*(0.933329 + 0.933314)
  if sampletag == 'cmass' and NorS == 1:
    vetof = 0.5*(0.906722 + 0.906815)

  if sampletag == 'lowz' and NorS == 0:
    vetof = 0.5*(0.933260+0.933202)
  if sampletag == 'lowz' and NorS == 1:
    vetof = 0.5*(0.906614+0.906454)

  tmp = mb.maskstats()
  assert tmp[NorS,0] == tmp[NorS,1] # weighted and unweighted areas are the same.
  maskarea = tmp[NorS,0]*vetof
  ##  make sure completeness is either 1 or 0.
  assert ((mb.weight == 1.) | (mb.weight == 0.)).all()
  ## old version.
  #return wgttot/maskarea
  return float(len(dd))/maskarea, wgttot/maskarea

def getdr12v4nbarcheck(nbarfname):
  """
  Computes target density per steradian as a check of the above calculation using the nbar files.
  This is less kludgy since nbar files have the veto fraction folded in.
  However, this removes stars from the nbar calculation so it should come out lower than 
  the true target distribution by the star fraction.
  """ 
  ifp = open(nbarfname,'r')
  line = ifp.readline()
  ifp.close()
  zcen,zlow,zhigh,nbar,shell_vol,total_gals = np.loadtxt(nbarfname,unpack=True,usecols=[0,1,2,3,5,6])
  ll = line.strip('\n').split(': ')[1].split(' ')
  for i in range(ll.count('')):
    ll.remove('')
  aeff = float(ll[0])*(np.pi/180.)**2
  #aeff = float(line.split(': ')[1].split(' ')[0])*(np.pi/180.)**2
  galtot = total_gals.sum()
  
  return galtot/aeff

def getnbarmean(nbarfname,zmin,zmax):
  zcen,zlow,zhigh,nbar,shell_vol,total_gals = np.loadtxt(nbarfname,unpack=True,usecols=[0,1,2,3,5,6])
  xx = np.where((zlow >= zmin*0.99999) & (zhigh <= zmax*1.00001))[0]
  print 'check it',zmin,zmax,zlow[xx].min(),zhigh[xx].max()
  return (nbar[xx]*total_gals[xx]).sum()/total_gals[xx].sum()


if __name__ == '__main__':

  nbardir = "/clusterfs/riemann/raid006/bosswork/boss/lss/dr12v4/"
  ddir = "/home/bareid/boss/bosslss/trunk/data/"
  
  for sample,zmin,zmax in zip(['lowz','cmass'],[0.2,0.43],[0.4,0.7]):
    for NorS,NS in zip([0,1],['N','S']):
      nbarstr1 = getdr12v4nbar(sampletag=sample,NorS=NorS)
      ## sanity check, this will be different because of the stars.
      nbarfname = nbardir + "nbar-"+sample+"-dr12v4-"+NS+"-Reid.dat"
      print nbarfname
      nbarstr2 = getdr12v4nbarcheck(nbarfname)
      nbar3d = getnbarmean(nbarfname,zmin,zmax)
      dd = fitsio.read(ddir+"collided-"+sample+"-dr12v4-"+NS+'-Reid.dat.fits')
      h,x = np.histogram(dd['IMATCH'],bins=np.arange(-0.5,10.5,1.0))
      fstar = h[4]/float(h[1:].sum())
  
      print '*'*80
      print 'results for ',sample,NS
      print 'best guess nbar',nbarstr1
      print 'nbar minus stars',nbarstr2
      print 'fraction low, stellar fraction',(nbarstr2-nbarstr1)/nbarstr1,fstar
      print 'nbar3d averaged over [%.2f, %.2f] = %e' % (zmin,zmax,nbar3d)
      print '*'*80