#----------------------------
# This script is used for resampling the data from Joyrad10.
# Author: Jose Dias Neto
# OPTIMIce Emmy-Noether Group
# Institute for Geophysics and Meteorology
# University of Cologne
#----------------------------


import resampleLib as rspl
from sys import argv
import pandas as pd
import numpy as np
import xarray as xr
import glob
import matplotlib.pyplot as plt
import os
import os.path

#----------------------------
# This block identifies the inputted argument.
# There are 3 possibilities for execution.
#
# 1) If no argument is passed, the code resample
# the data from the current day.
#
# e.g. --> python resampleJoyrad10.py
#
# 2) If the argument day_number is passed, the code
# executes for the current day - day_number.
#
# e.g. --> python resampleJoyrad10.py day_number:1
#
# If the current day is 20.03.2021, the example
# above will resample the data from 19.03.2021,
#
# 3) If the argument date is passed, the code
# resamples the data from the indicated day.
#
# e.g. --> python resampleJoyrad10.py date:20210101
#

try:

    mode, tmpDayIndex = argv[1].split(':')
    
    if mode == 'day_number':
        dayIndex = int(tmpDayIndex)

    if mode == 'date':
        date = pd.to_datetime(tmpDayIndex)

except:

    mode = 'day_number'
    dayIndex = 0
#
#----------------------------


#----------------------------
# This block identifies the date that will be resampled.
#

if mode == 'day_number':

    today = pd.datetime.today()
    deltaDay = pd.to_timedelta(dayIndex, 'D')

    startDay = pd.datetime(today.year, today.month, today.day) - deltaDay
    endDay = pd.datetime(today.year, today.month, today.day) - deltaDay


else:

    startDay = pd.datetime(date.year, date.month, date.day)
    endDay = pd.datetime(date.year, date.month, date.day)
 

#startDay = pd.to_datetime('20211003')
#endDay = pd.to_datetime('20211003')
print(startDay)

# list of days for resampling
daysToProc = pd.date_range(startDay, 
                           endDay,
                           freq='D')
#print(daysToProc)
#quit()
#
#----------------------------


#----------------------------
# This block defines the range reference grid.
# The reference grid is the same for all radars
#
beginRangeRef = 0 # starting height of the ref grid
endRangeRef = 12000 # ending height of the ref grid
rangeFreq = 36 # range resolution of the ref grid
rangeTolerance = 18 # tolerance for detecting the closest neighbour

rangeRef = np.arange(beginRangeRef, endRangeRef, rangeFreq)
#
#----------------------------

# tolerance for detecting closest neighbour (seconds)
timeTolerance = 1

# JOYRAD 10 input path
joyrad10Path = 'TEST/'
#'/archive/meteo/external-obs/juelich/joyrad10'

# resampled JOYRAD output path
joyrad10PathOutput = 'TEST/'
#joyrad10PathOutput = '/work/lvonterz/'

# Height offset, it is set to 0.32 here because the height
# of W-Band is the reference
rangeOffset = 0.32


#----------------------------
# This is the main processing block for
# for resampling the radar data
#
for date in daysToProc:
	# defining the final output path + name
	year = date.strftime('%Y')
	month = date.strftime('%m')
	day = date.strftime('%d')

	dateStr = year+month+day
	outPutFileName = ('/').join([joyrad10PathOutput,
		         'resampled',
		          dateStr+'_mon_joyrad10.nc'])
	#if not os.path.isfile(outPutFileName):
	# getting the time reference grid
	timeRef = rspl.getTimeRef(date)
	#print(timeRef)
	#print(rangeRed)
	#quit()
	
	# defining the input file name
	joyrad10FilePath = ('/').join([joyrad10Path,year, month, day,dateStr+'*.znc']) 

	# retrieving a list of files from the same day
	joyrad10FileList = sorted(glob.glob(joyrad10FilePath))
	#quit()
	if joyrad10FileList:
		#    joyrad10 = rspl.getVar(var, joyrad10, joyrad10FileList)
		try:
			joyrad10 = xr.open_mfdataset(joyrad10FileList,combine='by_coords')
			joyrad10 = joyrad10[['Zg','RMSg','VELg','SKWg']]
		except: 
			joyrad10 = xr.Dataset()
			for f in joyrad10FileList:
				try: 
					data = xr.open_dataset(f)
					joyrad10 = xr.merge([joyrad10,data[['Zg','RMSg','VELg','SKWg']]])
				except:
					print('cannot open ',f)
		_, index_time = np.unique(joyrad10['time'], return_index=True)
		joyrad10 = joyrad10.isel(time=index_time)
		joyrad10.time.attrs['units']='seconds since {0}'.format('1970-01-01 00:00:00 UTC')
		joyrad10 = xr.decode_cf(joyrad10)
		#

		# simplistic spurious data filtering
		joyrad10 = joyrad10.where(joyrad10['Zg']>0)

		# converting Zg to log units
		convert = ['Zg']
		for var in convert:
			joyrad10[var] = 10*np.log10(joyrad10[var])

		# correcting the range offset
		joyrad10.range.values = joyrad10.range.values + rangeOffset

		# resample along range
		joyrad10resRange = joyrad10.reindex({'range':rangeRef},method='nearest',tolerance=rangeTolerance)
		joyrad10resTime = joyrad10resRange.reindex({'time':timeRef},method='nearest',tolerance='4s')



		# saving the resampled data into a netCDF file
		print(outPutFileName)

		if os.path.exists(outPutFileName):
			os.remove(outPutFileName)
		encDic = {x: {"zlib": True} for x in joyrad10resTime}
		joyrad10resTime.to_netcdf(outPutFileName,encoding=encDic)
		joyrad10resTime.close()
	else:
		print('no files found ', dateStr)

	#else:
#		print('already processed file on ',dateStr)