#!/usr/local/bin/python2.7

#######################################################
# Author: Jonathan Case (ENSCO, Inc./SPoRT)
# Date: May 2016
# Name: imerg_hdf2grb.py
# Version: 1.0
#
# Details:
# This script is designed to read 30-min IMERG data
# files (either early, late, or final) and generate
# hourly GRIB2 files of accumulated precipitation
# estimates. The hourly GRIB2 files are for use in
# verifying WRF model precipitation using the SPoRT-MET
# scripting package.
#######################################################

# Modules:
import os,sys
import glob
import numpy as np
import array
import struct
from datetime import datetime
from datetime import timedelta
import time
import h5py

# scp to geo server?
doscp = 'y'
scpCmd = 'scp -q -o ConnectTimeout=5 -o ConnectionAttempts=1'
USER = 'casejl'
REMHOST = 'geo.nsstc.nasa.gov'

if (len(sys.argv) < 11):
    print "ERROR. Too few input arguments."
    print "Usage: imerg_hdf2grb.py yyyy mm dd hh lonLL lonUR latLL latUR region product"
    print "where yyyy mm dd hh is for the valid IMERG files,"
    print "and product is one of three IMERG products: early, late, or final."
    sys.exit()

script = sys.argv[0]
yyyy = sys.argv[1]
mm = sys.argv[2]
dd = sys.argv[3]
hh = sys.argv[4]
LonLL = sys.argv[5]
LonUR = sys.argv[6]
LatLL = sys.argv[7]
LatUR = sys.argv[8]
region = sys.argv[9]
product = sys.argv[10]

#print "Script arguments: ", str(sys.argv)
#sys.exit()

SRCDIR = '/usr/people/casejl/scripts/IMERG'

# When setting the valid date/time for 1-h APCP files, we need to go FORWARD 1 h so that the valid time is the END of the time window.
# This is the same convention as with STIV, MRMS, and CMORPH hourly GRIB2 files.
setdateDate = str(datetime.strptime(yyyy+mm+dd+hh, '%Y%m%d%H') + timedelta(hours=+1))  
vyyyy = setdateDate[0:4]
vmm = setdateDate[5:7]
vdd = setdateDate[8:10]
vhh = setdateDate[11:13]

if (product == 'early'):
  GRIDDATADIR = '/raid1/sport/data/IMERG/NRTE/'+yyyy+mm
  dataFile1 = '3B-HHR-E.MS.MRG.3IMERG.'+yyyy+mm+dd+'-S'+hh+'0000.HDF5'
  dataFile2 = '3B-HHR-E.MS.MRG.3IMERG.'+yyyy+mm+dd+'-S'+hh+'3000.HDF5'
  grbFile = GRIDDATADIR+'/imerge_'+region+'_'+vyyyy+vmm+vdd+vhh+'.grb2'
  REMDIR = '/data1/ftp/SPoRT/modeling/wrf/met/imerge'
elif (product == 'late'):
  GRIDDATADIR = '/raid1/sport/data/IMERG/NRTL/'+yyyy+mm
  dataFile1 = '3B-HHR-L.MS.MRG.3IMERG.'+yyyy+mm+dd+'-S'+hh+'0000.HDF5'
  dataFile2 = '3B-HHR-L.MS.MRG.3IMERG.'+yyyy+mm+dd+'-S'+hh+'3000.HDF5'
  grbFile = GRIDDATADIR+'/imergl_'+region+'_'+vyyyy+vmm+vdd+vhh+'.grb2'
  REMDIR = '/data1/ftp/SPoRT/modeling/wrf/met/imergl'
elif (product == 'final'):
  GRIDDATADIR = '/raid1/sport/data/IMERG/FINAL/'+yyyy+mm
  dataFile1 = '3B-HHR.MS.MRG.3IMERG.'+yyyy+mm+dd+'-S'+hh+'0000.HDF5'
  dataFile2 = '3B-HHR.MS.MRG.3IMERG.'+yyyy+mm+dd+'-S'+hh+'3000.HDF5'
  grbFile = GRIDDATADIR+'/imergf_'+region+'_'+vyyyy+vmm+vdd+vhh+'.grb2'
  REMDIR = '/data1/ftp/SPoRT/modeling/wrf/met/imergf'

#print "Dirs and Files: ",GRIDDATADIR,dataFile1,dataFile2,grbFile
#sys.exit()

num_lon = 3600
num_lat = 1800

# Open and read both half-hourly IMERG files
# Jayanthi's original code (below 2 lines)
#realprecip1=np.arange(num_lon*num_lat, dtype=np.float64)
#realprecip2=np.arange(num_lon*num_lat, dtype=np.float64)
# Why not initialize like other arrays? (below)
#realprecip1=np.zeros((num_lon*num_lat),dtype=np.float64).reshape((num_lat,num_lon),order='F')
#realprecip2=np.zeros((num_lon*num_lat),dtype=np.float64).reshape((num_lat,num_lon),order='F')
work1=np.zeros((num_lon*num_lat),dtype=np.float64).reshape((num_lon,num_lat),order='F')
work2=np.zeros((num_lon*num_lat),dtype=np.float64).reshape((num_lon,num_lat),order='F')
fnlpcp=np.zeros((num_lon*num_lat),dtype=np.float64).reshape((num_lon,num_lat),order='F')

# Process only if both half-hourly files are available.
if ( (os.path.isfile(GRIDDATADIR+'/'+dataFile1)) and (os.path.isfile(GRIDDATADIR+'/'+dataFile2)) ):
  file1 = h5py.File(GRIDDATADIR+'/'+dataFile1,"r")
  file2 = h5py.File(GRIDDATADIR+'/'+dataFile2,"r")
  name = '/Grid/precipitationCal'
  work1 = file1[name][:,:]
  work2 = file2[name][:,:]

  cond=(work1 < 0.0) | (work2 < 0.0)
# IMERG precip is in mm/h for both half-hourly blocks, so divide by 2.0 when summing two blocks.
  fnlpcp=np.where(cond,-9999.0,(work1+work2)/2.0)
  fnlpcp=np.reshape(fnlpcp,num_lon*num_lat,order='F')

#''' Binary output is written to unit 66 '''
  tmpoutfile = GRIDDATADIR+'/fort.66'
  fileobj = open(tmpoutfile,mode='wb')
  fileobj.write(struct.pack('<%df' % len(fnlpcp), *fnlpcp))
  fileobj.close()

#  sys.exit()

  globalGrb = GRIDDATADIR+'/imerg_global.grb2'
  os.system('/usr/people/casejl/bin/wgrib2lis '+SRCDIR+'/imergtemplate.grb2 -no_header -import_bin '+tmpoutfile+' -undefine_val -9.999000e+03 -set_date '+vyyyy+vmm+vdd+vhh+' -grib_out '+globalGrb)
  os.system('/usr/people/casejl/bin/wgrib2lis '+globalGrb+' -small_grib '+str(float(LonLL))+':'+str(float(LonUR))+' '+str(float(LatLL))+':'+str(float(LatUR))+' '+grbFile)

  os.unlink(tmpoutfile)
  os.unlink(globalGrb)

# J.Case (8/3/2018) -- correct logic if doscp is set to 'n'
#  if ( (doscp == 'y') and (os.path.isfile(grbFile)) ):
#    print "scp "+grbFile+" to "+REMHOST+":"+REMDIR
#    os.system(scpCmd+' '+grbFile+' '+USER+'@'+REMHOST+':'+REMDIR+'/')
#  else:
#    print "File "+grbFile+" does not exist. Not trying to scp...."
  if ( doscp == 'y' ):
    if ( os.path.isfile(grbFile) ):
      print "scp "+grbFile+" to "+REMHOST+":"+REMDIR
      os.system(scpCmd+' '+grbFile+' '+USER+'@'+REMHOST+':'+REMDIR+'/')
    else:
      print "File "+grbFile+" does not exist. Not trying to scp...."

else:
  print "Missing one or both IMERG files at hour "+hh

sys.exit()

