#+
#   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 scipy
import healpy as healpy
import os
from backend import *
from backend.cosmologyTools import *
import pickle

from netCDF4 import Dataset
NetCDFFile = Dataset
ncFloat = 'f8'

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

# -----------------------------------------------------------------------------
# returns the survey volume and scaled galaxy density for a given redshit slice
def getSurveyProps(sample):

  maskFile = sample.maskFile
  zmin = sample.zBoundary[0]
  zmax = sample.zBoundary[1]
  selFunMin = sample.zBoundary[0]
  selFunMax = sample.zBoundary[1]
  omegaM = sample.omegaM
  selectionFuncFile = sample.selFunFile
  useComoving = sample.useComoving

  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 volume normalization factors:
#   normalization used by zobov (assumes cubmic volume)
#   normalization used in galaxy survey volume for observations
def getVolNorm(sampleDir):

    with open(sampleDir+"/sample_info.dat", 'rb') as input:
      sample = pickle.load(input)

    infoFile = sampleDir+"/zobov_slice_"+sample.fullName+".par"
    File = NetCDFFile(infoFile, 'r')
    ranges = np.zeros((3,2))
    ranges[0][0] = getattr(File, 'range_x_min')
    ranges[0][1] = getattr(File, 'range_x_max')
    ranges[1][0] = getattr(File, 'range_y_min')
    ranges[1][1] = getattr(File, 'range_y_max')
    ranges[2][0] = getattr(File, 'range_z_min')
    ranges[2][1] = getattr(File, 'range_z_max')
    isObservation = getattr(File, 'is_observation')
    File.close()
    mul = np.zeros((3))
    mul[:] = ranges[:,1] - ranges[:,0]
    
    # getting the total number of tracers is an adventure depending
    # on which verison of VIDE we're pulling up
    maskIndexFile = sampleDir+"/mask_index.txt"
    totalPartFile = sampleDir+"/total_particles.txt"

    if os.path.exists(maskIndexFile):
      nTot = float(open(totalPartFile, "r").read())
      nGal = float(open(maskIndexFile, "r").read())
    else:
      nTot = float(open(totalPartFile, "r").read())
      nGal = nTot

    boxLen = mul

    boxVol = np.prod(boxLen)
    volNormZobov = nTot/boxVol
    volNormObs = volNormZobov

    if isObservation == 1:
      if os.access(sample.maskFile, os.F_OK):
        maskFile = sample.maskFile
      else:
        maskFile = sampleDir+"/"+os.path.basename(sample.maskFile)
        print("Using maskfile found in: " + maskFile)

      props = getSurveyProps(sample)
      surveyVol = props[0]
      volNormObs = nGal/surveyVol

    return volNormZobov, volNormObs


# -----------------------------------------------------------------------------
# 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 or redshift edge
def findEdgeGalaxies(galFile, maskFile, edgeGalFile, contourFile,
                     zmin, zmax, omegaM, useComoving,
                     meanPartSep, outputDir, log):

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

  log.write("  Reading contour map...\n")
  contourMap = healpy.read_map(contourFile)
  nside = healpy.get_nside(contourMap)
  npix = healpy.nside2npix(nside)

  # load in galaxies
  log.write("  Loading galaxies...\n")
  # TODO - WHY IS THIS FASTER THAN np.column_stack???
  galPos = np.genfromtxt(outputDir+"/galaxies.txt")
  with open(galFile, 'rb') as File:
   chk = np.fromfile(File, dtype=np.int32,count=1)
   Np = np.fromfile(File, dtype=np.int32,count=1)[0]
   chk = np.fromfile(File, dtype=np.int32,count=1)

   chk = np.fromfile(File, dtype=np.int32,count=1)
   x = np.fromfile(File, dtype=np.float32,count=Np)
   chk = np.fromfile(File, dtype=np.int32,count=1)

   chk = np.fromfile(File, dtype=np.int32,count=1)
   y = np.fromfile(File, dtype=np.float32,count=Np)
   chk = np.fromfile(File, dtype=np.int32,count=1)

   chk = np.fromfile(File, dtype=np.int32,count=1)
   z = np.fromfile(File, dtype=np.float32,count=Np)
   chk = np.fromfile(File, dtype=np.int32,count=1)

   chk = np.fromfile(File, dtype=np.int32,count=1)
   RA = np.fromfile(File, dtype=np.float32,count=Np)
   chk = np.fromfile(File, dtype=np.int32,count=1)

   chk = np.fromfile(File, dtype=np.int32,count=1)
   Dec = np.fromfile(File, dtype=np.float32,count=Np)
   chk = np.fromfile(File, dtype=np.int32,count=1)

   chk = np.fromfile(File, dtype=np.int32,count=1)
   redshift = np.fromfile(File, dtype=np.float32,count=Np)
   chk = np.fromfile(File, dtype=np.int32,count=1)

  #print(galPos.shape)
  #galPos = np.column_stack((x,y,z))
  #print(galPos.shape)

  log.write("  Building k-d tree...\n")
  flagList = np.zeros(len(galPos[:,0]))
  galTree = scipy.spatial.cKDTree(galPos)

  # flag galaxies near mask edges
  # using the "ray marching" algorithm: follow rays along lines of sight
  #   of all mask edges, flagging nearest neighbor galaxies as we go
 
  log.write("  Begin ray-marching...\n")
  rayDist = meanPartSep
  raySteps = np.arange(zmin, zmax, rayDist)
  log.write("  Will take %d steps per ray with %.2e distance between steps\n" % (len(raySteps), rayDist))
 
  contourPixels = np.nonzero(contourMap)[0]
  log.write("  We have %d rays to work with\n" % (len(contourPixels)))
  for pixel in contourPixels:
    #print("Working with pixel %d" % pixel)
    vec = healpy.pix2vec(nside,pixel)
    x = raySteps * vec[0]
    y = raySteps * vec[1]
    z = raySteps * vec[2]
    ray = np.array((x,y,z)).T
    
    dist, nearest = galTree.query(ray)
    flagList[nearest] = 1

  # flag galaxies near redsfhit boundaries
  # TODO - save time by only covering portion of sphere that has data

  log.write("  Flagging galaxies near redshift bounds...\n")
  sphereIndices = np.arange(len(contourMap))
  vec = healpy.pix2vec(nside, sphereIndices)
  vec = np.asarray(vec)
  maxEdge = zmax * vec
  maxEdge = maxEdge.T
  dist, nearest = galTree.query(maxEdge)
  flagList[nearest] = 2

  minEdge = zmin * vec
  minEdge = minEdge.T
  dist, nearest = galTree.query(minEdge)
  flagList[nearest] = 3

  log.write("  Flagging statistics:\n")
  log.write("    %d original galaxies\n" % len(flagList))
  log.write("    %d near edges\n" % len(flagList[flagList==1]))
  log.write("    %d near redshift bounds\n" % len(flagList[flagList==2]))
  log.write("    %d remaining\n" % len(flagList[flagList==0]))
  # output flag information
  log.write("  Saving galaxy flags to file...\n")
  np.savetxt(edgeGalFile, flagList, fmt="%d")

  # paint galaxy flags onto healpix map for diagnostics
  # TODO - drop this when done testing
  log.write("  Saving diagnostic maps to file...\n")
  flagMap = np.zeros(len(contourMap))
  justEdgeRA = RA[flagList == 1]
  justEdgeDec = Dec[flagList == 1]

  phi, theta = convertAngle(justEdgeRA, justEdgeDec)

  ipix = healpy.ang2pix(nside, theta, phi)
  for i in ipix:
    flagMap[i] += 1
  #np.put(flagMap, ipix, 1)

  healpy.write_map(outputDir+"/flagged_galaxies.fits", flagMap, 
                   overwrite=True,
                   dtype=np.dtype('float64'))


  allGalMap = np.zeros(len(contourMap))
  phi, theta = convertAngle(RA, Dec)
  
  ipix = healpy.ang2pix(nside, theta, phi)
  for i in ipix:
    allGalMap[i] += 1
  #np.put(allGalMap, ipix, 1)

  healpy.write_map(outputDir+"/all_galaxies.fits", allGalMap, 
                   overwrite=True,
                   dtype=np.dtype('float64'))

  return