#!/usr/bin/env python

# create surfclim.nc/surfinit.nc 
# Emanuel Dutra, November 2017 (adpated from CY43R1) 

from netCDF4 import Dataset,date2num
import time
import numpy as np
import gribapi as gb
import os
import datetime as dt


imiss=-9999
rmiss=1e20

ncVARS={}
ncVARS['ini']={}
ncVARS['clm']={}
ncVARS['veg']={}

ncVARS['veg']['lai_lv']={'dims':('time','x'),'long_name':'lai_lv','units':'-','dflt':-1}
ncVARS['veg']['lai_hv']={'dims':('time','x'),'long_name':'lai_hv','units':'-','dflt':-1}
ncVARS['veg']['tvl']=   {'dims':('time','x'),'long_name':'cvl','units':'-','dtype':'i4','dflt':-1}
ncVARS['veg']['tvh']=   {'dims':('time','x'),'long_name':'cvh','units':'-','dtype':'i4','dflt':-1}
ncVARS['veg']['cvl']=   {'dims':('time','x'),'long_name':'cvl','units':'-','dflt':-1}
ncVARS['veg']['cvh']=   {'dims':('time','x'),'long_name':'cvh','units':'-','dflt':-1}

ncVARS['ini']['SoilTemp']={'dims':('nlevs','x'),'long_name':'Soil Temperature','units':'K'}
ncVARS['ini']['SoilMoist']={'dims':('nlevs','x'),'long_name':'Soil Moisture','units':'m3/m3'}
ncVARS['ini']['iceTemp']={'dims':('nlevs','x'),'long_name':'Ice Temperature','units':'K','dflt':273}
ncVARS['ini']['SWE']={'dims':('x'),'long_name':'SWE','units':'kg m-2'}
ncVARS['ini']['snowdens']={'dims':('x'),'long_name':'Snow density','units':'kg m-3'}
ncVARS['ini']['SAlbedo']={'dims':('x'),'long_name':'Snow Albedo','units':'-'}
ncVARS['ini']['SnowT']={'dims':('x'),'long_name':'SnowT','units':'K'}
ncVARS['ini']['CanopInt']={'dims':('x'),'long_name':'CanopInt','units':'kg m-2','dflt':0.}
ncVARS['ini']['AvgSurfT']={'dims':('x'),'long_name':'AvgSurfT','units':'K'}
ncVARS['ini']['TLICE']={'dims':('x'),'long_name':'lake ice temperature','units':'K','dflt':273.}
ncVARS['ini']['TLMNW']={'dims':('x'),'long_name':'lake mean water temperature','units':'K','dflt':280.}
ncVARS['ini']['TLWML']={'dims':('x'),'long_name':'lake mixed layer temperature','units':'K','dflt':280.}
ncVARS['ini']['TLBOT']={'dims':('x'),'long_name':'lake bottom temperature','units':'K','dflt':280.}
ncVARS['ini']['TLSF']={'dims':('x'),'long_name':'lake temperature shape factor','units':'-','dflt':.65}
ncVARS['ini']['HLICE']={'dims':('x'),'long_name':'lake ice thickness','units':'m','dflt':0.}
ncVARS['ini']['HLML']={'dims':('x'),'long_name':'lake mixed layer thickness"','units':'K','dflt':5.}

ncVARS['clm']['Malbedo']={'dims':('month','x'),'long_name':'Monthly Albedo','units':'-','dflt':-1}
ncVARS['clm']['Mlail']=  {'dims':('month','x'),'long_name':'Monthly LAI low veg','units':'m2/m2','dflt':-1}
ncVARS['clm']['Mlaih']=  {'dims':('month','x'),'long_name':'Monthly LAI high veg','units':'m2/m2','dflt':-1}
ncVARS['clm']['Mask']=   {'dims':('x'),'long_name':'Mask','units':'-','comment':'1==run model','dtype':'i4','dflt':0}
ncVARS['clm']['z0m']=    {'dims':('x'),'long_name':'Surface roughness','units':'-','comment':'not used'}
ncVARS['clm']['lz0h']=   {'dims':('x'),'long_name':'log surface roughness length for heat','units':'-'}
ncVARS['clm']['landsea']={'dims':('x'),'long_name':'Land Sea Mask','units':'-'}
ncVARS['clm']['geopot']= {'dims':('x'),'long_name':'Geopotential (at the surface = orography)','units':'"m**2 s**-2"'}
ncVARS['clm']['cvl']=    {'dims':('x'),'long_name':'low vegetation cover','units':'-'}
ncVARS['clm']['cvh']=    {'dims':('x'),'long_name':'hight vegetation cover','units':'-'}
ncVARS['clm']['tvl']=    {'dims':('x'),'long_name':'type of low vegetation','units':'-','dtype':'i4'}
ncVARS['clm']['tvh']={'dims':('x'),'long_name':'type of hight vegetation','units':'-'}
ncVARS['clm']['sotype']={'dims':('x'),'long_name':'soil type','units':'-'}
ncVARS['clm']['sdor']={'dims':('x'),'long_name':'Standard deviation of orography','units':'-'}
ncVARS['clm']['sst']={'dims':('x'),'long_name':'sea surface temperature','units':'K','dflt':280.}
ncVARS['clm']['seaice']={'dims':('x'),'long_name':'sea ice mask','units':'-','dflt':0.}
ncVARS['clm']['LDEPTH']={'dims':('x'),'long_name':'Lake Depth','units':'m','dflt':10.}
ncVARS['clm']['CLAKE']={'dims':('x'),'long_name':'LAKE COVER','units':'-','dflt':0}
ncVARS['clm']['soaluvp']={'dims':('x'),'long_name':'Soil Albedo UV direct','units':'-'}
ncVARS['clm']['soaluvd']={'dims':('x'),'long_name':'Soil Albedo UV diffuse','units':'-'}
ncVARS['clm']['soalnid']={'dims':('x'),'long_name':'Soil Albedo NI diffuse','units':'-'}
ncVARS['clm']['soalnip']={'dims':('x'),'long_name':'Soil Albedo NI direct','units':'-'}

gb2NCveg={}
gb2NCveg['lai_lv']={'ncName':'lai_lv'}
gb2NCveg['lai_hv']={'ncName':'lai_hv'}
gb2NCveg['cvl']={'ncName':'cvl'}
gb2NCveg['cvh']={'ncName':'cvh'}
gb2NCveg['tvl']={'ncName':'tvl'}
gb2NCveg['tvh']={'ncName':'tvh'}

gb2NC={}
gb2NC['lsm']={'ncName':'landsea'}
gb2NC['cl']={'ncName':'CLAKE'}
gb2NC['sr']={'ncName':'z0m'}
gb2NC['lsrh']={'ncName':'lz0h'}
gb2NC['z']={'ncName':'geopot'}
gb2NC['slt']={'ncName':'sotype'}
gb2NC['dl']={'ncName':'LDEPTH'}
gb2NC['117.128']={'ncName':'soaluvp'}
gb2NC['118.128']={'ncName':'soaluvd'}
gb2NC['119.128']={'ncName':'soalnip'}
gb2NC['120.128']={'ncName':'soalnid'}

gb2NC['al']={'ncName':'Malbedo','klev':'month'}
gb2NC['lai_lv']={'ncName':'Mlail','klev':'month'}
gb2NC['lai_hv']={'ncName':'Mlaih','klev':'month'}

gb2NC['swvl1']={'ncName':'SoilMoist','klev':1,'lmin':0.01,'lmax':1.}
gb2NC['swvl2']={'ncName':'SoilMoist','klev':2,'lmin':0.01,'lmax':1.}
gb2NC['swvl3']={'ncName':'SoilMoist','klev':3,'lmin':0.01,'lmax':1.}
gb2NC['swvl4']={'ncName':'SoilMoist','klev':4,'lmin':0.01,'lmax':1.}
gb2NC['stl1']={'ncName':'SoilTemp','klev':1}
gb2NC['stl2']={'ncName':'SoilTemp','klev':2}
gb2NC['stl3']={'ncName':'SoilTemp','klev':3}
gb2NC['stl4']={'ncName':'SoilTemp','klev':4}
gb2NC['istl1']={'ncName':'iceTemp','klev':1}
gb2NC['istl2']={'ncName':'iceTemp','klev':2}
gb2NC['istl3']={'ncName':'iceTemp','klev':3}
gb2NC['istl4']={'ncName':'iceTemp','klev':4}
gb2NC['skt']={'ncName':'AvgSurfT'}
gb2NC['asn']={'ncName':'SAlbedo'}
gb2NC['rsn']={'ncName':'snowdens'}
gb2NC['sd']={'ncName':'SWE','mfac':1000.,'lmin':0.}
gb2NC['tsn']={'ncName':'SnowT'}
gb2NC['lmlt']={'ncName':'TLWML'}
gb2NC['lmld']={'ncName':'HLML'}
gb2NC['lblt']={'ncName':'TLBOT'}
gb2NC['ltlt']={'ncName':'TLMNW'}
gb2NC['lshf']={'ncName':'TLSF'}
gb2NC['lict']={'ncName':'TLICE'}
gb2NC['licd']={'ncName':'HLICE'}
gb2NC['src']={'ncName':'CanopInt','mfac':1000.,'lmin':0.}
#gb2NC['ci']={'ncName':'seaice'}


def gen_file(FOUT=None,FTYPE=None,ginfo=None,nlevs=4,nmonth=12):
  
  print "Generating:",FOUT
  os.system(' rm -f '+FOUT)
  
  ncfile = Dataset(FOUT,mode='w',format='NETCDF4_CLASSIC')
  ncfile.history = 'Created ' + time.ctime(time.time())

  ## dimensions
  ncfile.createDimension('x',ginfo['nptot'])
  ncfile.createDimension('nv',4)
  if FTYPE == 'ini':
    ncfile.createDimension('nlevs',nlevs)
  if FTYPE == 'clm':
    ncfile.createDimension('month',nmonth) 
  if FTYPE == 'veg':
    ncfile.createDimension('time',0)
  
 
  ## grid variables: 
  cvar=ncfile.createVariable('lat','f8',('x'))
  cvar.long_name = 'latitude'
  cvar.standard_name = 'latitude'
  cvar.units = 'degrees_north'
  cvar._CoordinateAxisType = "Lat"
  cvar.bounds = 'lat_bnds'
  cvar[:] = ginfo['lat']

  cvar=ncfile.createVariable('lon','f8',('x'))
  cvar.long_name = 'longitude'
  cvar.standard_name = 'longitude'
  cvar.units = 'degrees_east'
  cvar._CoordinateAxisType = "Lon"
  cvar.bounds = 'lon_bnds'
  cvar[:] = ginfo['lon']
  
  cvar=ncfile.createVariable('x','i4',('x',))
  cvar.long_name = 'Gaussian grid points 1-npoints'
  cvar.standard_name = 'Gaussian grid points 1-npoints'
  cvar.units = '-'
  cvar[:] = np.arange(ginfo['nptot'])+1

  cvar=ncfile.createVariable('lat_bnds','f8',('x','nv'),zlib=True,complevel=6)
  cvar.units = 'degrees_north'
  cvar[:] = ginfo['lat_bnds']
  
  cvar=ncfile.createVariable('lon_bnds','f8',('x','nv'),zlib=True,complevel=6)
  cvar.units = 'degrees_east'
  cvar[:] = ginfo['lon_bnds']

  cvar=ncfile.createVariable('cell_area','f8',('x'),zlib=True,complevel=6)
  cvar.long_name = 'area of grid cell'
  cvar.standard_name = 'area'
  cvar.units = 'm2'
  cvar.coordinates = 'lat lon'
  cvar[:] = ginfo['area']
  
  if FTYPE == 'clm':
    cvar=ncfile.createVariable('month','f4',('month',))
    cvar.long_name = 'month'
    cvar.units = 'months since 2000-01-01'
    ncfile.variables['month'][:]=[0,1,2,3,4,5,6,7,8,9,10,11]
  
  if FTYPE == 'veg':
    cvar=ncfile.createVariable('time','f4',('time',))
    cvar.long_name = 'time'
    cvar.units = 'seconds since XXX'  
  
  for cv,val in ncVARS[FTYPE].iteritems():
    try:
      dd = val['dtype']
    except:
      dd = 'f4'
    #print cv,dd,val
    cvar=ncfile.createVariable(cv,dd,val['dims'],zlib=True,complevel=6)
    for kk in val.keys():
      if kk in ['dims','dtype','dflt']: continue
      setattr(cvar,kk,val[kk])
    try:
      f = val['dflt']
      cvar[:] = f
    except:
      pass


  ncfile.close()
  
def grib2nc(FNC,FGRB,gb2NC=gb2NC,FTYPE=None):
  print "reading:",FGRB
  print "writing to:",FNC
  nc = Dataset(FNC,'r+')
  grb = open(FGRB)
  
  ifld=0
  while 1:
    gid = gb.grib_new_from_file(grb)
    if gid is None: break
    ifld = ifld + 1
    sname = gb.grib_get(gid,'shortName')
    param = gb.grib_get(gid,'param')
    if sname == "~":
      sname=str(param)
    kmm   = gb.grib_get(gid,'month')
    
    klev=None
    mfac=1
    lmin=-1.e20
    lmax=1.e20
    ncName=None
    if gb2NC.has_key(sname):
      ncName = gb2NC[sname]['ncName']
      if gb2NC[sname].has_key('klev'):
        klev=gb2NC[sname]['klev']
      else:
        klev=None
      if gb2NC[sname].has_key('mfac'):
        mfac=gb2NC[sname]['mfac']
      if gb2NC[sname].has_key('lmin'):
        lmin=gb2NC[sname]['lmin']
      if gb2NC[sname].has_key('lmax'):
        mfac=gb2NC[sname]['lmax']
    else:
      if sname in nc.variables.keys():
        ncName = sname
        
    if FTYPE == 'veg':
      xdd = gb.grib_get(gid,'validityDate')
      xdT = gb.grib_get(gid,'validityTime')/100.
      xdateR = (dt.datetime.strptime('%8i'%(xdd),'%Y%m%d')+
                                dt.timedelta(hours=xdT) )
      if ifld==1:
        xdate=xdateR
        klevt=0
        #nctime = "seconds since "+xdate.strftime('%Y-%m-%d %H:%M:%S')
        nctime = "seconds since "+xdate.isoformat(' ')
        nc.variables['time'].units = nctime 
        nc.variables['time'][klevt] = date2num(xdate,nctime)
        print 'FIRST:',klevt,xdate
      else:
        if xdate != xdateR:
          xdate = xdateR
          klevt=klevt+1
          nc.variables['time'][klevt] = date2num(xdate,nctime)
          print 'OTHERS:',klevt,xdate
      klev=klevt
        
    print ifld,sname,ncName,klev,mfac,kmm
    
    if ncName is not None:    
      xdata = gb.grib_get_array(gid,'values')
      xdata = np.maximum(lmin,np.minimum(lmax,xdata * mfac) )
      if FTYPE == 'veg':
        nc.variables[ncName][klevt,:] = xdata
      elif klev is None:
        nc.variables[ncName][:] = xdata 
      elif klev == 'month':
        nc.variables[ncName][kmm-1,:] = xdata 
      else:
        nc.variables[ncName][klev-1,:] = xdata 
        
        
    gb.grib_release(gid)
    
 
  nc.close()
  grb.close()
  
def get_mask(FIN,LLAND,LLAKE,LOCEAN,BBOX):
  nc = Dataset(FIN,'r')
  lsm = nc.variables['landsea'][:]
  clake = nc.variables['CLAKE'][:]
  lat = nc.variables['lat'][:]
  lon = nc.variables['lon'][:]
  lon[lon>180.] = lon[lon>180.] - 360. 
  
  
  LDLAND = lsm >= 0.5 
  LDLAKE = clake >= 0.5*(1-lsm)
  LDLAKE_RESOL = clake >= 0.5
  LDOCEAN = (lsm < 0.5 ) & (~LDLAKE)
  LDREG = ( (lon >= BBOX[0]) & (lon <= BBOX[1]) &
            (lat >= BBOX[2]) & (lat <= BBOX[3]) )
  
  LDMASK = LDREG & (LDLAND*LLAND | LDLAKE*LLAKE | LDOCEAN*LOCEAN)
  print "NPTOT:",lsm.size
  print "NPLAND:",np.sum(LDLAND)
  print "NPLAKE:",np.sum(LDLAKE)
  print "NPLAKE_RESOL:",np.sum(LDLAKE_RESOL)
  print "NPOCEAN:",np.sum(LDOCEAN)
  print "NPREG:",np.sum(LDREG)
  print "NPMASK:",np.sum(LDMASK)
  
  return LDMASK
  
def cut_mask(mask,FIN,FOUT):
  
  NPTOT = np.sum(mask)
  print "Generating:",FOUT,NPTOT 
  os.system(' rm -f '+FOUT)
  ncOUT = Dataset(FOUT,mode='w',format='NETCDF4_CLASSIC')
  ncOUT.history = 'Created ' + time.ctime(time.time()) + ' based on '+FIN
  
  ncIN = Dataset(FIN,'r')
  
  
  ## copy dimensions
  for cdim in ncIN.dimensions.keys():
    if ncIN.dimensions[cdim].isunlimited():
        dlen=None
    else:
        dlen = len(ncIN.dimensions[cdim])
    if cdim == 'x':
      dlen = NPTOT
    ncOUT.createDimension(cdim,dlen)
  
  for vname, varin in ncIN.variables.iteritems():
    vtype = varin.dtype
    vdim = varin.dimensions
    vfill = None
    zlib = True
    complevel=6
    if hasattr(varin,'_FillValue'):
        vfill = varin._FillValue
    #if hasattr(varin,'_Shuffle'):
        #zlib = True
    #if hasattr(varin,'_DeflateLevel'):
        #complevel = getattr(varin,'_DeflateLevel')
    print vname,vtype,vdim,vfill,zlib,complevel
    cdfvar = ncOUT.createVariable(vname,vtype,vdim,fill_value=vfill,
                  zlib=zlib,complevel=complevel)
   
    ## copy attributes 
    for catt in varin.ncattrs():
      if catt == '_FillValue': continue
      ccat = getattr(varin,catt)
      setattr(cdfvar,catt,ccat)
      
    xIN = ncIN.variables[vname][:]
    ndim = xIN.ndim
    if ndim == mask.ndim and cdfvar.size == NPTOT :
      cdfvar[:] = xIN[mask]
    elif vname in  ['lat_bnds','lon_bnds']:
      cdfvar[:] = xIN[mask,:]
    elif ndim == mask.ndim + 1 :
      cdfvar[:] = xIN[:,mask]
    elif (ndim == mask.ndim -1) and (cdfvar.size == NPTOT):
      cdfvar[:] = xIN[mask.squeeze()]
    if vname == 'time':
      cdfvar[:] = xIN[:]

  
  ncOUT.close()
  ncIN.close()
  

#FCLIMgrb='/scratch/rd/need/osm_devel_ini_clm/clim.grb'
#FINITgrb='/scratch/rd/need/osm_devel_ini_clm/init.grb'

### Generate netcdf files (empty)
#ginfo = gg.get_grib_grid(FCLIMgrb,retbounds=True) 
#gen_file(FOUT='surfclim_all',FTYPE='clm',ginfo=ginfo)
#gen_file(FOUT='soilinit_all',FTYPE='ini',ginfo=ginfo)

### Convert grib to netcdf 
#grib2nc(FNC='surfclim_all',FGRB=FCLIMgrb)
#grib2nc(FNC='soilinit_all',FGRB=FINITgrb)

#### Defaults defining points 
#LLAND=True
#LLAKE=True
#LOCEAN=False
#BBOX=[-190,190,-91,91]

### find mask, i.e. point where we'll ran the model 
#mask = get_mask('surfclim_all',LLAND,LLAKE,LOCEAN,BBOX)

### Create final files cutting down mask 
#cut_mask(mask=mask,FIN='surfclim_all',FOUT='surfclim')
#cut_mask(mask=mask,FIN='soilinit_all',FOUT='soilinit')
### finally set Mask to 1 
#nc = Dataset('surfclim','r+')
#nc.variables['Mask'][:] = 1
#nc.close()

