#!/usr/bin/env python

# create forcing files
# Emanuel Dutra, November 2017 

from netCDF4 import Dataset,date2num
import time
import numpy as np
import eccodes as ec 
import os
import datetime as dt


vinfo={}
vinfo['Wind']={'lname':'Wind speed','units':'m/s'}
vinfo['Tair']={'lname':'Temperature','units':'K'}
vinfo['Qair']={'lname':'Specific humidity','units':'kg/kg'}
vinfo['PSurf']={'lname':'Surface Pressure','units':'Pa'}
vinfo['SWdown']={'lname':'Downward shortwave radiation','units':'W/m2'}
vinfo['LWdown']={'lname':'Downward longwave radiation','units':'W/m2'}
vinfo['Rainf']={'lname':'Rainfall','units':'kg/m2/s'}
vinfo['Snowf']={'lname':'Snowfall','units':'kg/m2/s'}
vinfo['Ctpf']={'lname':'Convective total precipitation fraction','units':'-'}
vinfo['Rhair']={'lname':'Relative humidity','units':'-'}

def read_grid(fname):
  nc = Dataset(fname,'r')
  gpoints = nc.variables['x'][:]  #indexed from 1---npoints 
  lat = nc.variables['lat'][:]
  lon = nc.variables['lon'][:]
  nc.close()
  return gpoints,lat,lon

def gen_output(cvar,tunits,gpoints,lat,lon):
  fout=cvar+'.nc'
  print "Generating",fout
  try:
    os.remove(fout)
  except:
    pass
  nc = Dataset(fout,'w',format='NETCDF4')
  nc.createDimension('x',len(gpoints))
  nc.createDimension('time',0)
  
  ncvar = nc.createVariable('x','i4',('x',))
  ncvar.long_name = 'Gaussian grid points 1-npoints'
  ncvar.standard_name = 'Gaussian grid points 1-npoints'
  ncvar.units = '-'
  ncvar[:] = gpoints
  
  ncvar = nc.createVariable('lat','f4',('x',))
  ncvar.long_name = 'latitude'
  ncvar.standard_name = 'latitude'
  ncvar.units = 'degrees_north'
  ncvar._CoordinateAxisType = "Lat"
  ncvar[:] = lat

  ncvar=nc.createVariable('lon','f4',('x',))
  ncvar.long_name = 'longitude'
  ncvar.standard_name = 'longitude'
  ncvar.units = 'degrees_east'
  ncvar._CoordinateAxisType = "Lon"
  ncvar[:] = lon
  
  ncvar=nc.createVariable('time','f8',('time',))
  ncvar.long_name = 'time'
  ncvar.standard_name = 'time'
  ncvar.units = tunits
  ncvar._CoordinateAxisType = "time"
  ncvar[0] = 0
  
  
  ncvar=nc.createVariable(cvar,'f4',('time','x'),zlib=True,complevel=6)
  ncvar.long_name = vinfo[cvar]['lname']
  ncvar.standard_name = vinfo[cvar]['lname']
  ncvar.units = vinfo[cvar]['units']
  ncvar._CoordinateAxisType = "Time Lat Lon"
  ncvar[0,:] = -1
  return nc

def read_grib(FGRB,Sname,gpoints):
  print "Reading:",Sname,FGRB
  grb = open(FGRB)
  ifld=0
  xdata=[]
  xtime=[]
  while True:
    gid = ec.codes_grib_new_from_file(grb)
    if gid is None: break
    sname = ec.codes_get(gid,'shortName')
    if Sname == sname:
      xdata.append(ec.codes_get_values(gid))
      xdd = ec.codes_get(gid,'validityDate')
      xdT = ec.codes_get(gid,'validityTime')
      xtime.append(dt.datetime.strptime(str(xdd)+str(xdT).zfill(4),'%Y%m%d%H%M'))
      ifld=ifld+1
    ec.codes_release(gid)
  grb.close()
  xtime = np.array(xtime)
  xdata = np.array(xdata)
  print "Read",xdata.shape
  xdata = xdata[:,gpoints-1]
  return xdata,xtime

def calc_rh_t_td(t,td):
  return np.maximum(PRHMIN,np.minimum(PRHMAX,np.exp(-RLVTT/(RV*t*td)*(t-td))))

def calc_qair(rh,t,p):
  pes=calc_PES(t)
  return rh*pes / ( p*(RETV+1.)-rh*pes*RETV)

def calc_rh_p_q_t(p,q,t):
  pes = calc_PES(t)
  return np.maximum(PRHMIN,np.minimum(PRHMAX,(p*q*(RETV+1.)) / ((1.+RETV*q)*pes) ) )

def calc_PES(Tair):
  
  # ratio between water and ice : Alpha in eq(7.73) for IFS Document Cy38R1e staturation vapor
  FOEALFA = np.minimum(1.0,( (np.maximum(
                       RTICE,np.minimum(RTWAT,Tair))-RTICE) /(RTWAT-RTICE))**2.0)
  #pressure with Teten's formula, eq. 7.15 for IFS doc CY38R1
  FOEEWM = R2ES *(FOEALFA*np.exp(R3LES*(Tair-RTT)/(Tair-R4LES))+
           (1.0-FOEALFA)*np.exp(R3IES*(Tair-RTT)/(Tair-R4IES)))
  return FOEEWM


def process(SURFCLIM,INPUTGRB,FTYPE,tunits):
  # 1. Get grid info
  gpoints,lat,lon = read_grid(SURFCLIM)
  # 2. Create output variable 

  if FTYPE== 'SFC':
    t2,xtime = read_grib(INPUTGRB,'2t',gpoints)
    d2,xtime = read_grib(INPUTGRB,'2d',gpoints)
    psurf,xtime = read_grib(INPUTGRB,'sp',gpoints)
    rh = calc_rh_t_td(t2,d2)
    q2 = calc_qair(rh,t2,psurf)
    CVAR='Tair'
    nc=gen_output(CVAR,tunits,gpoints,lat,lon)
    nc.variables[CVAR][:]=t2
    nc.variables['time'][:]=date2num(xtime,tunits)
    nc.close()
    CVAR='Qair'
    nc=gen_output(CVAR,tunits,gpoints,lat,lon)
    nc.variables[CVAR][:]=q2
    nc.variables['time'][:]=date2num(xtime,tunits)
    nc.close()
    ## just for testing
    #CVAR='Rhair'
    #nc=gen_output(CVAR,tunits)
    #nc.variables[CVAR][:]=rh
    #nc.variables['time'][:]=date2num(xtime,tunits)
    #nc.close()
    
  elif FTYPE == 'ML':
    CVAR='Tair'
    xdata,xtime = read_grib(INPUTGRB,'t',gpoints)
    nc=gen_output(CVAR,tunits,gpoints,lat,lon)
    nc.variables[CVAR][:]=xdata
    nc.variables['time'][:]=date2num(xtime,tunits)
    nc.close()
    CVAR='Qair'
    xdata,xtime = read_grib(INPUTGRB,'q',gpoints)
    nc=gen_output(CVAR,tunits,gpoints,lat,lon)
    nc.variables[CVAR][:]=xdata
    nc.variables['time'][:]=date2num(xtime,tunits)
    nc.close()

  CVAR='PSurf'
  xdata,xtime = read_grib(INPUTGRB,'sp',gpoints)
  nc=gen_output(CVAR,tunits,gpoints,lat,lon)
  nc.variables[CVAR][:]=xdata
  nc.variables['time'][:]=date2num(xtime,tunits)
  nc.close()

  CVAR='Wind'
  if FTYPE== 'SFC':
    x1,xtime = read_grib(INPUTGRB,'10u',gpoints)
    x2,xtime = read_grib(INPUTGRB,'10v',gpoints)
  elif FTYPE == 'ML':
    x1,xtime = read_grib(INPUTGRB,'u',gpoints)
    x2,xtime = read_grib(INPUTGRB,'v',gpoints)
  xdata = np.sqrt(x1*x1+x2*x2)
  nc=gen_output(CVAR,tunits,gpoints,lat,lon)
  nc.variables[CVAR][:]=xdata
  nc.variables['time'][:]=date2num(xtime,tunits)
  nc.close()

  CVAR='SWdown'
  xdata,xtime = read_grib(INPUTGRB,'ssrd',gpoints)
  nc=gen_output(CVAR,tunits,gpoints,lat,lon)
  nc.variables[CVAR][:]=np.maximum(0,xdata/float((xtime[1]-xtime[0]).seconds))
  nc.variables['time'][:]=date2num(xtime,tunits)
  nc.close()

  CVAR='LWdown'
  xdata,xtime = read_grib(INPUTGRB,'strd',gpoints)
  nc=gen_output(CVAR,tunits,gpoints,lat,lon)
  nc.variables[CVAR][:]=np.maximum(0,xdata/float((xtime[1]-xtime[0]).seconds))
  nc.variables['time'][:]=date2num(xtime,tunits)
  nc.close()

  ## Precip special:
  lsp,xtime = read_grib(INPUTGRB,'lsp',gpoints)
  cp,xtime = read_grib(INPUTGRB,'cp',gpoints)
  sf,xtime = read_grib(INPUTGRB,'sf',gpoints)
  rainf = lsp+cp - sf
  ctpf = np.where(lsp+cp > 0,cp / (lsp+cp),0)

  CVAR='Rainf'
  nc=gen_output(CVAR,tunits,gpoints,lat,lon)
  nc.variables[CVAR][:]=np.maximum(0,rainf/float((xtime[1]-xtime[0]).seconds))*1000.
  nc.variables['time'][:]=date2num(xtime,tunits)
  nc.close()

  CVAR='Snowf'
  nc=gen_output(CVAR,tunits,gpoints,lat,lon)
  nc.variables[CVAR][:]=np.maximum(0,sf/float((xtime[1]-xtime[0]).seconds))*1000.
  nc.variables['time'][:]=date2num(xtime,tunits)
  nc.close()

  CVAR='Ctpf'
  nc=gen_output(CVAR,tunits,gpoints,lat,lon)
  nc.variables[CVAR][:]=np.minimum(1,np.maximum(0,ctpf))
  nc.variables['time'][:]=date2num(xtime,tunits)
  nc.close()
  
  return


#CONSTANTS:
RLVTT=2.5008E+6
RG=9.80665            # Gravitational Acceleration [m/s2]
RKBOL=1.380658E-23    # Boltzmann constant (k) [J/K]
RNAVO=6.0221367E+23   # Avogadro constant (N_A) [1/mol]
R=RNAVO*RKBOL         # Perfect gas constant (= 8.314511211948600)
RMD=28.9644           # Dry air mass
RMV=18.0153           # Vapour  mass
RD=1000.0*R/RMD       # Dry air cst. (= 287.0596736665907 J/kg/K)
RV=1000.0*R/RMV       # Vapour  cst. (= 461.5249933083879 J/kg/K)
RETV=RV/RD-1.0        # Rv/Rd-1 (= 0.608)
RTT=273.16            # Temperature of water fusion at normal pressure
RTWAT=RTT             # Criteria for temperature with respect to liquid water
RTICE=RTT-23.0        # Criteria for temperature with respect to ice
R2ES=611.21           # Constant in Teten's formula: eq(7.5)
R3LES=17.5020         # Mixing Ratio Over Liquid Water in Teten's formula: eq(7.5)
R4LES=32.190         # Constant in Teten's formula for liquid water: eq(7.5)
R3IES=22.5870        # Mixing Ratio Over Ice in Teten's formula: eq(7.5)
R4IES=-0.70          # Constant in Teten's formula for ice: eq(7.5)
PRHMAX=0.999          # max relative humidity
PRHMIN=0.001          # Min relative humidiy


#SURFCLIM="/disk1/data/work/ec-earth/test/osm/inidata/fix/surfclim"
#INPUTGRB="tmp_gg.grb"
#FTYPE="SFC"
#tunits='seconds since 1990-01-01 00:00:00'

#process(SURFCLIM,INPUTGRB,FTYPE,tunits)



  
