#!/usr/bin/env python

# Script to convert netcdf to grib
# Emanuel Dutra, November 2017 (adapted from CY43r1)

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


class variable:
  def __init__(self,pID=None,table=None,ncfile=None,ncvar=None,nlev=0,mfact=1.,isflx=False):
     self.pID = pID
     self.table = table
     self.ncfile = ncfile
     self.ncvar = ncvar
     self.nlev = nlev
     self.mfact = mfact
     self.isflx = isflx

ptable={}
## AN type 
ptable['39.128']=variable(39,128,'o_gg.nc','SoilMoist',1,1./(0.07*1000.))
ptable['40.128'] =variable(40,128,'o_gg.nc','SoilMoist',2,1./(0.21*1000.))
ptable['41.128'] =variable(41,128,'o_gg.nc','SoilMoist',3,1./(0.72*1000.))
ptable['42.128'] =variable(42,128,'o_gg.nc','SoilMoist',4,1./(1.89*1000.))
ptable['139.128'] =variable(139,128,'o_gg.nc','SoilTemp',1,)
ptable['170.128'] =variable(170,128,'o_gg.nc','SoilTemp',2,)
ptable['183.128'] =variable(183,128,'o_gg.nc','SoilTemp',3,)
ptable['236.128'] =variable(236,128,'o_gg.nc','SoilTemp',4,)
ptable['235.128'] =variable(235,128,'o_gg.nc','AvgSurfT',)
ptable['238.128'] =variable(238,128,'o_gg.nc','SnowT')
ptable['141.128'] =variable(141,128,'o_gg.nc','SWE',mfact=0.001)
ptable['33.128'] =variable(33,128,'o_gg.nc','snowdens')
ptable['32.128'] =variable(32,128,'o_gg.nc','SAlbedo')
ptable['198.128'] =variable(198,128,'o_gg.nc','CanopInt',mfact=0.001)
ptable['8.228'] =variable(8,228,'o_gg.nc','TLWML')
ptable['9.228'] =variable(9,228,'o_gg.nc','HLML')
ptable['10.228'] =variable(10,228,'o_gg.nc','TLBOT')
ptable['11.228'] =variable(11,228,'o_gg.nc','TLMNW')
ptable['12.228'] =variable(12,228,'o_gg.nc','TLSF')
ptable['13.228'] =variable(13,228,'o_gg.nc','TLICE')
ptable['14.228'] =variable(14,228,'o_gg.nc','HLICE')

## FC type 
ptable['167.128'] =variable(167,128,'o_d2m.nc','T2m')
ptable['243.128'] =variable(243,128,'o_sus.nc','Albedo')
ptable['168.128'] =variable(168,128,'o_d2m.nc','D2m')
ptable['66.128'] =variable(66,128,'o_vty.nc','lai',nlev=1)
ptable['67.128'] =variable(67,128,'o_vty.nc','lai',nlev=2)
ptable['260038'] =variable(260038,-1,'o_cld.nc','SnowFrac',mfact=100)
ptable['3066'] =variable(3066,-1,'o_cld.nc','SnowDepth')

## flx type 
ptable['176.128'] =variable(176,128,'o_efl.nc','SWnet',isflx=True)
ptable['177.128'] =variable(177,128,'o_efl.nc','LWnet',isflx=True)
ptable['169.128'] =variable(169,128,'o_efl.nc','SWdown',isflx=True)
ptable['175.128'] =variable(175,128,'o_efl.nc','LWdown',isflx=True)
ptable['147.128'] =variable(147,128,'o_efl.nc','Qle',isflx=True)
ptable['146.128'] =variable(146,128,'o_efl.nc','Qh',isflx=True)

ptable['80.228'] =variable(80,228,'o_co2.nc','CO2flux',isflx=True)
ptable['81.228'] =variable(81,228,'o_co2.nc','Ag',isflx=True)
ptable['82.228'] =variable(82,228,'o_co2.nc','Reco',isflx=True)

ptable['144.128'] =variable(144,128,'o_wat.nc','Snowf',isflx=True,mfact=0.001)
ptable['182.128'] =variable(182,128,'o_wat.nc','Evap',isflx=True,mfact=0.001)
ptable['45.128'] =variable(45,128,'o_wat.nc','Qsm',isflx=True,mfact=0.001)
ptable['228.128'] =variable(228,128,'o_wat.nc','Rainf',isflx=True,mfact=0.001)
ptable['8.128'] =variable(8,128,'o_wat.nc','Qs',isflx=True,mfact=-0.001)
ptable['9.128'] =variable(9,128,'o_wat.nc','Qsb',isflx=True,mfact=-0.001)
ptable['205.128'] =variable(205,128,'o_wat.nc','Qs',isflx=True,mfact=-0.001)
ptable['44.128'] =variable(44,128,'o_eva.nc','SubSnow',isflx=True,mfact=0.001)

##  extra evap. fluxes using eccodes 2.3 
ptable['228100'] =variable(228100,-1,'o_eva.nc','ECanop',isflx=True,mfact=0.001)
ptable['228101'] =variable(228101,-1,'o_eva.nc','TVeg',isflx=True,mfact=0.001)
ptable['228102'] =variable(228102,-1,'o_eva.nc','ESoil',isflx=True,mfact=0.001)
ptable['228103'] =variable(228103,-1,'o_eva.nc','EWater',isflx=True,mfact=0.001)
##  extra evap. fluxes using eccodes 2.3 
ptable['100.228'] =variable(100,228,'o_eva.nc','ECanop',isflx=True,mfact=0.001)
ptable['101.228'] =variable(101,228,'o_eva.nc','TVeg',isflx=True,mfact=0.001)
ptable['102.228'] =variable(102,228,'o_eva.nc','ESoil',isflx=True,mfact=0.001)
ptable['103.228'] =variable(103,228,'o_eva.nc','EWater',isflx=True,mfact=0.001)

## mulit-layer fields 
# soil 
ptable['228093']=variable(228093,-1,'o_gg.nc','SoilMoist',-1,)
ptable['228139'] =variable(228139,-1,'o_gg.nc','SoilTemp',-1,)
# snow 
ptable['228141']=variable(228141,-1,'o_gg.nc','SWEML',-1,)
ptable['33']=variable(33,-1,'o_gg.nc','snowdensML',-1,)
ptable['190141']=variable(190141,-1,'o_gg.nc','slwML',-1,)
ptable['201203']=variable(201203,-1,'o_gg.nc','SnowTML',-1,)



def write2grb(CFSURF,CFTEMP,CFBASE,ZMISS,FCTYPE,FGOUT,PARAMS,LACUM=True,LRESET=-1,DATE0=None,TIME0=0,NSTEP=1):


  ## open mask 
  ncCLM = Dataset(CFSURF,'r')
  JPOINTS = ncCLM.variables['x'][:]
  ncCLM.close()
  NPCDF=JPOINTS.size

  ## open template 
  fgrb = open(CFTEMP)
  gid = gb.grib_new_from_file(fgrb)
  NPTOT = gb.grib_get(gid,'numberOfPoints')

  ## open output
  fout = open(FGOUT,'w')


  for kk in PARAMS:
    vinfo = ptable[kk]
    nc = Dataset(CFBASE+vinfo.ncfile,'r')
    cvar = nc.variables[vinfo.ncvar]
    ndims=len(cvar.dimensions)
    NPCDF1 = len(nc.dimensions[cvar.dimensions[-1]])
    NT1 = len(nc.dimensions[cvar.dimensions[0]])
    NT0=0
    NLEV1=0
    if ndims == 3:
      NLEV1 = len(nc.dimensions[cvar.dimensions[1]])
    print 'CVAR,CODE,NPCDF,NLEV1,NT1:',vinfo.ncvar,vinfo.pID,NPCDF1,NLEV1,NT1

    
   
    ZTIME = num2date(nc.variables['time'][:],getattr(nc.variables['time'],'units'))
    if (len(ZTIME) > 1 ):
      DSTEP=(ZTIME[1]-ZTIME[0]).total_seconds()
    else:
      DSTEP=nc.variables['time'][0]  

    if vinfo.ncfile == 'o_gg.nc' : NT0 = 1   # special case of o_gg skip first 
    gidC = gb.grib_clone(gid)
    gedition = gb.grib_get(gidC,'editionNumber')
    if vinfo.table == -1 : 
      gb.grib_set(gidC,'editionNumber',2)
      gedition = 2
    if gedition == 2:
      gb.grib_set(gidC,'param',vinfo.pID) 
    else:
      gb.grib_set(gidC,'indicatorOfParameter',vinfo.pID)
      gb.grib_set(gidC,'table2Version',vinfo.table)
    
    gb.grib_set(gidC,'dataType','FC') # hardcoded 
    gb.grib_set(gidC,'bitmapPresent',1)
    gb.grib_set(gidC,'missingValue',ZMISS) 
    gb.grib_set(gidC,'bitsPerValue',16) 
    if vinfo.pID == 141:
      gb.grib_set(gidC,'bitsPerValue',24)
    
    if FCTYPE == 'FC':
      if DATE0 == -1:
        if (len(ZTIME) > 1 ):
          DDATE0 = int(ZTIME[0].strftime('%Y%m%d'))
        else:
          DDATE0 = int((ZTIME[0]-dt.timedelta(0,DSTEP)).strftime('%Y%m%d'))
      else:
        DDATE0 = DATE0   
      print DDATE0,TIME0
      date00 = dt.datetime.strptime(str(DDATE0)+str(TIME0),'%Y%m%d%H')
      
     ## load all data
     
    if vinfo.nlev >0: 
      ZIN = cvar[:,vinfo.nlev-1,:]
    else:
      ZIN = cvar[:]
    
    if vinfo.pID == 228 and vinfo.table == 128 :
      # special case for total precip:
      ZIN = ZIN + nc.variables['Snowf'][:]
    if vinfo.pID == 205 and vinfo.table == 128 :
      # special case for Total runoff:
      ZIN = ZIN + nc.variables['Qsb'][:]  
      
      
    ZOUT =  np.zeros((NPTOT))+ZMISS 
    mfact=vinfo.mfact
    if vinfo.isflx:
      mfact = vinfo.mfact * DSTEP
      if LACUM:
        ZOUT[JPOINTS-1]=0.
    print mfact
    ## Write grid output 
    for jt in range(NT0,NT1,NSTEP):
      
      #ZOUT =  np.zeros((NPTOT))+ZMISS
      
      if LACUM and  vinfo.isflx :
        ZOUT[JPOINTS-1] = ZOUT[JPOINTS-1] + ZIN[jt,:]*mfact
      else:
        ZOUT[JPOINTS-1] = ZIN[jt,:]*mfact
    
      gb.grib_set_values(gidC,ZOUT)
      
      step=0
      dataTime=0
      if FCTYPE == 'AN':
        dataDate=int(ZTIME[jt].strftime("%Y%m%d"))
        dataTime=100*int(ZTIME[jt].strftime("%H"))
      elif FCTYPE == 'FC':
        dataDate=DDATE0
        dataTime=TIME0*100
        step=(ZTIME[jt]-date00).total_seconds()/3600.
        
      print jt,ZTIME[jt],dataDate,dataTime,step
      gb.grib_set(gidC,'dataDate',dataDate)
      gb.grib_set(gidC,'dataTime',dataTime)
      gb.grib_set(gidC,'step',step)
      
      
      gb.grib_write(gidC,fout)
      
      if LACUM and  vinfo.isflx and FCTYPE == 'FC' :
        if LRESET == -1 :
          ZOUT[JPOINTS-1]=0.
          print 'RESET FLX'
        else:
          if np.mod(step,LRESET)==0:
            ZOUT[JPOINTS-1]=0.
            print 'RESET FLX'
            
  gb.grib_release(gidC)
  fgrb.close()
  fout.close()
  nc.close()
  

#write2grb(CFSURF="/home/emanuel/data/work/ec-earth/test/osm/inidata/fix/surfclim",
          #CFTEMP="/home/emanuel/data/work/ec-earth/test/osm/inidata/fix/init.grb",
          #CFBASE="/home/emanuel/data/work/ec-earth/test/osm/outdata/19900101/",
          #ZMISS=9999,
          #FCTYPE='AN',
          #FGOUT='out.grb',
          #PARAMS=['39.128','40.128','41.128','42.128',
                  #'139.128','170.128','183.128','236.128',
                  #'235.128','238.128','141.128','33.128','32.128',
                  #'167.128','243.128','168.128','66.128','67.128',
                  #'176.128','177.128','147.128','146.128','182.128',
                  #'205.128','169.128','175.128','144.128'],
         #LACUM=False,LRESET=-1,DATE0=None,TIME0=0)
##grib_copy -B "date:i asc, time asc" out.grb out1.grb


