## generic functions to handle gaussian reduced gribs
## Emanuel Dutra, November 2017

import eccodes as ec 
import numpy as np 

PLANET_RADIUS=6371000.  # default radius for area calculations

def get_reduced_lats_weights(N,FLOC=''):
  """ 
  Compute the latitudes and weigths of a gaussian grid with N latitudes:
  lats,weights = get_reduced_lats_weights(N) 
  
  Using:
  Legendre quadrature points on [-1,1]. Returns the roots of Nth order Legendre polynomial P_N(x) and   
  weights (w) to use in Gaussian Quadrature over [-1,1] with weighting function 1.
  
  Adapted from  Anna Agusti-Panareda
  """
  #import scipy.special
  #sinlat,weight = scipy.special.orthogonal.p_roots(N)
  import os.path
  fname=FLOC+'gauss_%i.npy'%(N/2)
  if os.path.isfile(fname):
    OUT = np.load(fname)
    lat = OUT[:,0]
    weight = OUT[:,1]
  else:
    OUT=np.zeros((N,2))
    OUT[:,0],OUT[:,1] = np.polynomial.legendre.leggauss(N)
  
    lat=(np.arcsin(OUT[:,0])*180.0/np.pi)[::-1] # invert to be between 90: -90
    weight = OUT[:,1][::-1]
  return lat,weight

def compute_lats_weights(FLOC='/home/rd/need/python/pyutils/'):
  NN =np.array([16,32,48,80,96,128,160,200,256,320,400,512,640,800,1024,1280,1600])
  
  for N in NN:
    print N
    OUT = np.zeros((N*2,2))
    OUT[:,0],OUT[:,1]  = get_reduced_lats_weights(N*2,FLOC='/xxx')


def get_grib_grid(fin,verbose=True,retbounds=False):
  """ 
  Get grib grid information 
  grid_info = get_grib_grid(fin,verbose)
  fin - grib file (will read the first field only)
  verbose - if true prints some information
  grid_info: dictionary with the output:
    grid_info['nptot'] = npG # total number of grid points 
    grid_info['nlat'] = nlatG # number of latitude lines
    grid_info['nlon'] = 0 # to identify we'r gaussian !
    grid_info['nplon'] = nplonG # number longitudes in each latitude 
    grid_info['lat'] = latG # global lat
    grid_info['lon'] = lonG # global lon
    grid_info['nplonC'] = nplonGC  # cumulative number of points in each latitude line 
    grid_info['latr'] = latrG  # gaussian regular latitude
    grid_info['lonr'] = lonrG  # gaussian regular longitude
    grid_info['dlon'] = dlon  # longitude increment in each latitude line
    grid_info['weights'] =   # gaussian weigths
  """
  if verbose:
    print 'get_grib_grid: reading file',fin
  f = open(fin)
  gid = ec.codes_grib_new_from_file(f)
  
  ## information from grib file 
  nlatG = ec.codes_get(gid,'Nj') # number of latitudes
  nplonG = np.empty(nlatG,dtype='i4') 
  nplonG = ec.codes_get_array(gid,'pl')  # number of point in each latitude
  ec.codes_release(gid)
  f.close()
  ## derived
  nplonGC = np.cumsum(nplonG,dtype='i4')
  npG = np.sum(nplonG,dtype='i4')
  lonr = np.linspace(0.,360.,np.max(nplonG),endpoint=False)
  dlonG = 360./nplonG
  ## gen gaussian lat and weights
  latr,weights = get_reduced_lats_weights(nlatG)
  
  latG = np.empty(npG,dtype='f8')
  lonG = np.empty(npG,dtype='f8')
  istart=0
  for ilat in xrange(nlatG):
    iend = istart+nplonG[ilat]
    latG[istart:iend] = latr[ilat]
    lonG[istart:iend] = np.linspace(0.,360.,nplonG[ilat],endpoint=False)
    istart = iend


  grid_info = {}
  grid_info['nptot'] = npG # total number of grid points 
  grid_info['nlat'] = nlatG # number of latitude lines
  grid_info['nlon'] = 0 # to identify we'r gaussian !
  grid_info['nplon'] = nplonG # number longitudes in each latitude 
  grid_info['lat'] = latG # global lat
  grid_info['lon'] = lonG # global lon
  grid_info['nplonC'] = nplonGC  # cumulative number of points in each latitude line 
  grid_info['latr'] = latr  # gaussian regular latitude
  grid_info['lonr'] = lonr  # gaussian regular longitude
  grid_info['dlon'] = dlonG  # longitude increment in each latitude line
  grid_info['gweights'] = weights  # longitude increment in each latitude line
  
  
  if retbounds:
    xbounds=np.zeros((npG,4),dtype='f4')  # west,east,east,west
    ybounds=np.zeros((npG,4),dtype='f4')  # south,south,north,north
    #
    #       3 -----  2
    #         \   \
    #         \   \
    #       0 -----  1
    #

    ## compute lat edges
    latedges = np.zeros((nlatG+1),dtype='f8')
    latedges[0] = 90.
    latedges[1:-1] = np.arcsin(1-np.cumsum(weights[:-1]))*180/np.pi
    latedges[-1] = -90.
    ## filll the arrays 
    istart=0
    for ilat in xrange(nlatG):
      iend = istart+nplonG[ilat]
      ybounds[istart:iend,0]=latedges[ilat+1]
      ybounds[istart:iend,1]=latedges[ilat+1]
      ybounds[istart:iend,2]=latedges[ilat]
      ybounds[istart:iend,3]=latedges[ilat]
      xbounds[istart:iend,0]=lonG[istart:iend]-dlonG[ilat]*.5
      xbounds[istart:iend,1]=lonG[istart:iend]+dlonG[ilat]*.5
      xbounds[istart:iend,2]=xbounds[istart:iend,1]
      xbounds[istart:iend,3]=xbounds[istart:iend,0]
      istart = iend

    grid_info['lon_bnds'] = xbounds  # longitude bounds
    grid_info['lat_bnds'] = ybounds  # longitude bounds

    ## area computations
  area= np.empty(npG,dtype='f8')
  #area = np.empty(npG,dtype='f8')
  istart=0
  for ilat in xrange(nlatG):
    iend = istart+nplonG[ilat]
    #xarea = 0.5*np.abs(np.sin(np.deg2rad(ybounds[istart,0]))-
                       #np.sin(np.deg2rad(ybounds[istart,2])) )/nplonG[ilat]
    xarea = 4*np.pi*PLANET_RADIUS**2*weights[ilat] / (nplonG[ilat]*np.sum(weights))
    area[istart:iend] = xarea
    #area1[istart:iend] = xarea1 
    istart = iend

  grid_info['area'] = area  # longitude increment in each latitude line
  #grid_info['area'] = area  # longitude increment in each latitude line

  if verbose:
    print "nptot,nlat,nlon:",' ',grid_info['nptot'],grid_info['nlat'],grid_info['nlon']
    print 'get_grib_grid: finished'
  return grid_info
