#+
#   VIDE -- Void IDentification and Examination
#   Copyright (C) 2010-2014 Guilhem Lavaux
#   Copyright (C) 2011-2014 P. M. Sutter
#
#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation; version 2 of the License.
# 
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License along
#   with this program; if not, write to the Free Software Foundation, Inc.,
#   51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#+
# a suite of functions to compute expansion rates, angular diameter 
# distances, and expected void stretching

import numpy as np
import healpy as healpy
import os
from backend import *
from backend.cosmologyTools import *

__all__=['getSurveyProps', 'getNside', 'figureOutMask', 'findEdgeGalaxies']

# -----------------------------------------------------------------------------
# returns the volume and galaxy density for a given redshit slice
def getSurveyProps(maskFile, zmin, zmax, selFunMin, selFunMax, portion, 
                   omegaM, selectionFuncFile=None, useComoving=False):

  #LIGHT_SPEED = 299792.458

  mask = healpy.read_map(maskFile)
  area = (1.*np.size(np.where(mask > 0)) / np.size(mask)) * 4.*np.pi

  if useComoving:
    zmin = LIGHT_SPEED/100.*comovingDistance(zmin, Om=omegaM)
    zmax = LIGHT_SPEED/100.*comovingDistance(zmax, Om=omegaM)
  else:
    zmin = zmin * LIGHT_SPEED/100.
    zmax = zmax * LIGHT_SPEED/100. 

  volume = area * (zmax**3  - zmin**3) / 3

  if selectionFuncFile == None:
    nbar = 1.0

  elif not os.access(selectionFuncFile, os.F_OK):
    print("   Warning, selection function file %s not found, using default of uniform selection." % selectionFuncFile)
    nbar = 1.0

  else:
    selfunc = np.genfromtxt(selectionFuncFile)
    selfunc = np.array(selfunc)
    selfunc[:,0] = selfunc[:,0]/100.
    selfuncUnity = selfunc
    selfuncUnity[:,1] = 1.0
    selfuncMin = selfunc[0,0]
    selfuncMax = selfunc[-1,0]
    selfuncDx = selfunc[1,0] - selfunc[0,0]
    selfuncN = np.size(selfunc[:,0])
  
    selFunMin = max(selFunMin, selfuncMin)
    selFunMax = min(selFunMax, selfuncMax)
  
  
    def f(z): return selfunc[np.ceil((z-selfuncMin)/selfuncDx), 1]*z**2 
    def fTotal(z): return selfuncUnity[np.ceil((z-selfuncMin)/selfuncDx), 1]*z**2 
  
    zrange = np.linspace(selFunMin, selFunMax)
  
    nbar = scipy.integrate.quad(f, selFunMin, selFunMax)
    nbar = nbar[0]
  
    ntotal = scipy.integrate.quad(fTotal, 0.0, max(selfuncUnity[:,0]))
    ntotal = ntotal[0]
  
    nbar = ntotal / area / nbar


  return (volume, nbar)


# -----------------------------------------------------------------------------
# returns the nside resolution from the given maskfile
def getNside(maskFile):

  mask = healpy.read_map(maskFile)
  return healpy.get_nside(mask)


# -----------------------------------------------------------------------------
# helper function to convert RA,dec to phi,theta
def convertAngle(RA, Dec):

  phi   = np.pi/180.*RA
  theta = Dec*np.pi/180.
  theta = np.pi/2. - Dec*np.pi/180.

  return (phi, theta)

# -----------------------------------------------------------------------------
# computes the mask from a given galaxy datafile and writes it to a file
def figureOutMask(galFile, nside, outMaskFile):

  npix = healpy.nside2npix(nside)
  mask = np.zeros((npix))

  for line in open(galFile):
    line = line.split()
    RA  = float(line[3])
    Dec = float(line[4])
    z   = float(line[5])

    phi, theta = convertAngle(RA, Dec)

    pix = healpy.ang2pix(nside, theta, phi)
    mask[pix] = 1.
 
  healpy.write_map(outMaskFile, mask, overwrite=True, dtype=np.dtype('float64'))

  return mask

# -----------------------------------------------------------------------------
# figures out which galaxies live on a mask edge, and also writes the edge
#   map to an auxillary file
def findEdgeGalaxies(galFile, maskFile, edgeGalFile, edgeMaskFile,
                     zmin, zmax, omegaM, useComoving, boundaryWidth):

  if useComoving:
    zmin = comovingDistance(zmin, Om=omegaM)
    zmax = comovingDistance(zmax, Om=omegaM)
    #zmin = LIGHT_SPEED/100.*comovingDistance(zmin, Om=omegaM)
    #zmax = LIGHT_SPEED/100.*comovingDistance(zmax, Om=omegaM)
  #else:
  #  zmin *= LIGHT_SPEED/100.
  #  zmax *= LIGHT_SPEED/100.


  mask = healpy.read_map(maskFile)
  nside = healpy.get_nside(mask)
  npix = healpy.nside2npix(nside)
  edgeMask = np.zeros((npix))

  edgeFile = open(edgeGalFile, "w")

  for line in open(galFile):
    line = line.split()
    RA  = float(line[3])
    Dec = float(line[4])
    z   = float(line[5])

    if useComoving:
      z = comovingDistance(z/LIGHT_SPEED, Om=omegaM)
    else:
      z *= LIGHT_SPEED/100.

    phi, theta = convertAngle(RA, Dec)

    # check the mask edges
    ipix = healpy.ang2pix(nside, theta, phi)
    neighbors = healpy.get_all_neighbours(nside, ipix)
    isOnMaskEdge = any(mask[p] == 0 for p in neighbors)
  
    # check the redshift boundaries
    zbuffer = (zmax-zmin)*boundaryWidth
    isOnHighZEdge = (z >= zmax-zbuffer)
    isOnLowZEdge  = (z <= zmin+zbuffer)

    if isOnMaskEdge:
      edgeFile.write("1\n")
      edgeMask[ipix] = 1
    elif isOnHighZEdge:
      edgeFile.write("2\n")
    elif isOnLowZEdge:
      edgeFile.write("3\n")
    else:
      edgeFile.write("0\n")

  edgeFile.close()
  healpy.write_map(edgeMaskFile, edgeMask, overwrite=True,
                   dtype=np.dtype('float64'))

  return