from headers import *

##################################################################################
#  Mathematical functions



def W3d_sth(x):
   """Fourier transform of a 3d spherical top hat window function.
   Use x = k*R as input,
   where R is the tophat radius and k the wave vector.
   Input and output are dimensionless.
   """
   if x < 1.e-3:  # for small x, replace by expansion, for numerical stability
      f = 1. - 0.1* x**2 + 0.00357143* x**4
   else:
      f = (3./(x**3)) * ( np.sin(x) - x * np.cos(x) )
   return f


def dW3d_sth(x):
   """Derivative of the FT of the top hat.
   Input and output are dimensionless.
   """
   f = 3. * (3. * x * np.cos(x) - 3. * np.sin(x) + (x**2) * np.sin(x)) / (x**4)
   return f


def W2d_cth(x):
   """FT of a 2d circular top hat window function.
   Input and output are dimensionless.
   """
   return 2.*special.jn(1, x) / x

def W1d_th(x):
   """FT of a 1d tophat
   normalized to unity at k=0 (ie real space integral is 1)
   Input and output are dimensionless.
   """
   return sinc(x/2.)
   
def Si(x):
   return special.sici(x)[0]

def Ci(x):
   return special.sici(x)[1]

def sinc(x):
   return special.sph_jn(0, x)[0][0]

def j0(x):
   """relevant for isotropic Fourier transform in 2d
   """
   return special.jn(0, x)


def i0(x):
   """Modified Bessel function of the first kind
   """
   return special.iv(0, x)


##################################################################################
# formatting numbers

def intExpForm(input, round=None):
   """
   clean scientific notation for file names
   removes trailing decimal point if not needed
   """
   a = '%e' % np.float(input)
   # mantissa: remove trailing zeros
   mantissa = a.split('e')[0].rstrip('0')
   # round if desired
   if round is not None:
      intPart, decPart = mantissa.split('.')
      mantissa = intPart + '.' + decPart[:round]
   # mantissa: remove trailing zeros
   # then remove dot if no decimal digits
   mantissa = mantissa.rstrip('0').rstrip('.')
   # exponent: remove + sign if there, and leading zeros
   exponent = np.int(a.split('e')[1])
   exponent = np.str(exponent)
   if exponent=='0':
      return mantissa
   else:
      return mantissa + 'e' + exponent



def floatExpForm(input, round=None):
   """same as intExpForm, except always leaves the decimal point
   round: number of digits to keep after the dot
   """
   a = '%e' % np.float(input)
   # mantissa: remove trailing zeros
   mantissa = a.split('e')[0].rstrip('0')
   # round if desired
   if round is not None:
      intPart, decPart = mantissa.split('.')
      mantissa = intPart + '.' + decPart[:round]
   # exponent: remove + sign if there, and leading zeros
   exponent = np.int(a.split('e')[1])
   exponent = np.str(exponent)
   if exponent=='0':
      return mantissa
   else:
      return mantissa + 'e' + exponent



##################################################################################
# Matrix inversion for ill-conditioned matrices, with SVD

def invertMatrixSvdTruncated(matrix, epsilon=1.e-5, keepLow=True):
   '''Invert a matrix by inverting its SVD,
   and setting to zero the singular values that are too small/large.
   epsilon sets the tolerance for discarding singular values.
   keepLow=True: for inverting a cov matrix, want to keep the modes with lowest variance.
   keepLow=False: for inverting a Fisher maytrix, want to keep the modes with highest Fisher information.
   '''
   # Perform SVD on matrix
   U, s, Vh = scipy.linalg.svd(matrix)
   # invert the singular values
   sInv = 1./s
   # remove the super poorly constrained modes, that lead to numerical instabilities
   if keepLow:
      sInvMax = np.max(sInv)
      sInv[sInv<=sInvMax*epsilon] = 0.
   else:
      sInvMin = np.min(sInv)
      sInv[sInv>=sInvMin/epsilon] = 0.
   # create diagonal matrix
   sInv = np.diag(sInv)
   # invert the hermitian matrices
   V = np.conj(Vh.transpose())
   Uh = np.conj(U.transpose())
   # generate the inverse
   result = np.dot(V, np.dot(sInv, Uh))
   return result


##################################################################################
# Generating ell bins with constant number of modes


def generateEllBins(lMin, lMax, nL, fsky=1.):
   '''Generates nL bins between lMin and lMax,
   such that the number of 2d modes in each bin is identical.
   Returns the bin centers, the bin edges, the bin widths, and the number of modes per bin.
   '''
   # area in ell space between lMin and l,
   # normalized to 1 when l=lMax
   farea = lambda l: (l**2 - lMin**2) / (lMax**2 - lMin**2)

   # find the bin edges,
   # such that each bin has equal number of modes
   Le = np.zeros(nL+1)
   for iL in range(nL+1):
      f = lambda l: farea(l) - float(iL) / nL
      Le[iL] = optimize.brentq(f , lMin, lMax)
   
   # use the average ell in the bin, weighted by number of modes, as bin center
   Lc =  2./3. * (Le[1:]**3 - Le[:-1]**3) / (Le[1:]**2 - Le[:-1]**2)
   # bin spacing
   dL = Le[1:] - Le[:-1]
   # Number of modes
   lF = 2.*np.pi / np.sqrt(4. * np.pi * fsky)
   nModesTotal = np.pi * (lMax**2 - lMin**2) / lF**2
   Nmodes = nModesTotal / nL * np.ones(nL)   
   
   return Lc, dL, Nmodes, Le


##################################################################################
#  Extract non-mask data vector and cov matrix

def extractMaskedMat(cov, mask=None, I=None):
   '''cov: large matrix
   mask: 1d array, 0 if unmasked, anything else if masked
   I: indices of the large matrix to keep, pre-masking
   '''
   # convert mask to 0 and 1
   mask = mask.astype(bool)
   # extract indices of interest, if needed
   if I is not None:
      mask = mask[I]
      J = np.ix_(I, I)
      cov = cov[J]
   if mask is not None:
      # nb of unmasked rows
      nNew = np.int(np.sum(1 - mask))
      # mask cov matrix
      mask = np.diag(mask)
      cov = ma.masked_array(cov, mask=mask)
      cov = ma.mask_rowcols(cov)
      # extract the non-masked elements
      cov = cov.compressed().reshape((nNew, nNew))
   return cov

def extractMaskedVec(vec, mask=None, I=None):
   '''vec: large vector
   mask: 1d array, 0 if unmasked, anything else if masked
   I: indices of the large vector to keep, pre-masking
   '''
   # convert mask to 0 and 1
   mask = mask.astype(bool)
   # extract indices of interest, if needed
   if I is not None:
      mask = mask[I]
      vec = vec[I]
   if mask is not None:
      # mask vec matrix
      vec = ma.masked_array(vec, mask=mask)
      # extract the non-masked elements
      vec = vec.compressed()
   return vec


def splitBins(x, nBins):
   '''Give the nBins+1 edges of the nBins bins,
   such that there is an equal number of elements of x
   in each bin.
   '''
   nBins = np.int(nBins)
   # sort the x values
   xSorted = np.sort(x)
   # corresponding probability
   proba = 1. * np.arange(len(x)) / (len(x) - 1)
   # interpolate the CDF
   # scipy is smart: even if twice the same x appears,
   # it will know to give the interpolating function the right jump
   cdf = interp1d(xSorted, proba, kind='linear', bounds_error=False, fill_value=(0., 1.))

   # fill an array with the z-bounds of the bins
   binEdges = np.zeros(nBins+1)
   binEdges[0] = np.min(x)
   binEdges[-1] = np.max(x)
   for iBin in range(nBins-1):
      # force equal number of objects per bins
      f = lambda xMax: cdf(xMax) - (iBin+1.)/nBins
      binEdges[iBin+1] = optimize.brentq(f , binEdges[0], binEdges[-1])
   return binEdges



