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


import matplotlib
matplotlib.use('Agg')

import matplotlib.pyplot as plt
from sys import argv
import pandas as pd
import xarray as xr
import numpy as np
import plottingLib as plib
import resampleLib as rsp
import os

# figures output path
#pathOut = '/home/l/L.Terzi/radar_quicklooks/plots'
pathOut = 'TEST/'

# path to emptyData
emptyDataPath = '/scratch/l/L.Terzi/campaign_aux_data/tripex-pol/auxPlotData/noData.nc'

# path to joyrad10 resampled data
dataPath10 = '/scratch/l/L.Terzi/campaign_aux_data/tripex-pol/joyrad10/resampled'

# path to joyrad35 resampled data
dataPath35 = '/scratch/l/L.Terzi/campaign_aux_data/tripex-pol/joyrad35/resampled'

# path to grarad94 resampled data
# dataPath94 = '/data/optimice/tripex-pol/grarad94/resampled'

# path to mirac94 resampled data
#dataPath94 = '/data/optimice/tripex-pol/mirac94/resampled'

# path to wband_scan resampled data
dataPath94 = '/scratch/l/L.Terzi/campaign_aux_data/tripex-pol-scan/wband_scan/resampled'

#----------------------------
# 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 tripex_pol_plots.py
#
# 2) If the argument day_number is passed, the code
# executes for the current day - day_number.
#
# e.g. --> python tripex_pol_plots.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 tripex_pol_plots.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 
#
#----------------------------


#----------------------------
#
try:

    remove, state = argv[2].split(':')

except:

    state = False

print('remove .nc files: {0}'.format(state))
#----------------------------


#----------------------------
# This block identifies the date that will be resampled.
#
if mode == 'day_number':

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

    startDay = today - deltaDay
    endDay = today - deltaDay


else:

    startDay = pd.to_datetime(date)
    endDay = pd.to_datetime(date)

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


#----------------------------
# This is the main processing block for
# for plotting the resampled data

for day in daysToProc:

    # date identification
    date = day.strftime('%Y%m%d')
    print('plotting: {0}'.format(date))

    # trying to oppen the resampled joyrad10 data
    try:

        fileName10 = date+'_mon_joyrad10.nc'
        filePath10 = ('/').join([dataPath10, fileName10])
        data10 = xr.open_dataset(filePath10)
        
    # reading an empty dataset in case joyrad10 does not exist
    except:

        tmpData = xr.open_dataset(emptyDataPath)
        data10 = tmpData.copy()
        data10['time']=rsp.getTimeRef(day)
        tmpData.close()
        print("couldn't open joyrad10 file at ", filePath10)
        
    # trying to oppen the resampled joyrad35 data
    try:

        fileName35 = date+'_mon_joyrad35.nc'
        filePath35 = ('/').join([dataPath35, fileName35])
        data35 = xr.open_dataset(filePath35)

    # creating an empty dataset in case joyrad35 does not exist
    except:

        tmpData = xr.open_dataset(emptyDataPath)
        data35 = tmpData.copy()
        #data35['time']=rsp.getTimeRef(day)
        tmpData.close()

    # trying to oppen the grarad94 orher 94 GHz radar
    try:

        fileName94 = date+'_ZEN_moments_wband_scan.nc'
        filePath94 = ('/').join([dataPath94, fileName94])
        data94 = xr.open_dataset(filePath94)
        data94 = data94.rename({'Ze':'Zg','MDV':'VELg','WIDTH':'RMSg','sLDR':'sLDR_w','SK':'SKWg'})
    # creating an empty dataset in case grarad does not exist
    except:
        fileName94 = date+'_ZEN_moments_wband_scan.nc'
        filePath94 = ('/').join([dataPath94, fileName94])
        print('couldnt find data94 file at ',filePath94)
        
        tmpData = xr.open_dataset(emptyDataPath)
        data94 = tmpData.copy()
        #data94['time']=rsp.getTimeRef(day)
        tmpData.close()

    #---------------------------------------
    # Use this block in case an Ze offset
    # correction is needed
    #
    #data10['Zg'] = data10['Zg'] + offsetX
    #data35['Zg'] = data35['Zg'] + offsetKa
    #data94['Zg'] = data94['Zg'] + offsetW
    #---------------------------------------

    #-- plot all moments
    variables = {'Zg':{'vmax':25, 'vmin':-35,'units':'[dB]'},
                 'VELg':{'vmax':0, 'vmin':-3,'units':r'[ms$^{-1}$]'},
                 'RMSg':{'vmax':0, 'vmin':1,'units':r'[ms$^{-1}$]'},
                 'SKWg':{'vmax':1, 'vmin':-1,'units':r'[]'},
                }

    # creating the triple panels plot
    
    for var in variables.keys():

        plib.plotVar(data35[var], data94[var],
                variables[var]['vmax'], variables[var]['vmin'],
                pathOut, date, var,variables[var]['units'],CEL=False,data10=data10[var])
        print(var,' plotted ZEN')
    # creating LDR plot, not right now because we don't have Ka-Band here
    
    #plib.plotLDR(data35, -20, -35, pathOut, date, 'LDR_ka')
    #print('ZEN LDR plotted')
    
    # defining the color range and the name of the
    # differences used by the plotting function
    variables = {'Zg':{'vmax':20, 'vmin':-5, 'name':'DWR','units':'[dB]'},
                 'VELg':{'vmax':0.3, 'vmin':-0.3, 'name':'DDV','units':r'[ms$^{-1}$]'},
                 'RMSg':{'vmax':0.3, 'vmin':-0.3, 'name':'DSW','units':r'[ms$^{-1}$]'}
                }

    # creating the difference plots
    for var in variables.keys():

        diff1035 = data10[var].astype('float32') - data35[var].astype('float32')
        diff3594 = data35[var].astype('float32') - data94[var].astype('float32')

        diff1035.attrs['long_name']=variables[var]['name']+'-XKa'
        diff3594.attrs['long_name']=variables[var]['name']+'-KaW'

        plib.plotDiffVar(diff3594,
                    variables[var]['vmax'], variables[var]['vmin'],
                    pathOut, date , variables[var]['name'],variables[var]['units'],CEL=False,diff1035=diff1035)
    print('plotted difference variable ZEN')
    # closing all files
    data10.close()
    data35.close()
    data94.close()