#!/export/hpc-lw-hmtdev1/strojniak/miniconda3/bin/python

#######
# Sarah Python /export/hpc-lw-hmtdev1/strojniak/miniconda3/bin/python
#Jimmy python /export/wpc-lw-hmtdev4/jcorreia/miniconda3/bin/python
#
# Updated : Sarah Trojniak Apr 14 2022
#   new directory paths and names
###################

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt 
import numpy as np
import os
from datetime import datetime, timedelta
import xarray as xr
import cfgrib
import pygrib
from pyproj import Proj, transform
import cartopy.crs as ccrs
import cartopy
from matplotlib.colors import LinearSegmentedColormap
import matplotlib.colors as mcolors
import argparse, copy
import subprocess, pickle
import json
from matplotlib.path import Path
import scipy.ndimage as ndimage
from skimage.morphology import dilation,disk
from scipy.spatial import cKDTree
import sys
#func_loc = '/export/wpc-lw-hmtdev1/strojniak/ffair/code/'
#sys.path.insert(0, func_loc)
from hmt_functions_ensv3a import *
#np.set_printoptions(threshold=sys.maxsize)

debug = False
add_snow = False

# python plot_exp.py -m 'NSSL' -c '00' -t '0' -s 'True'
# init argparse
parser = argparse.ArgumentParser()
parser.add_argument('-m', '--model', help='Model', type=str)
parser.add_argument('-c', '--cycle', help='Cycle', type=str)
parser.add_argument('-t', '--time', help='time', type=str)
parser.add_argument('-n', '--snow', help='snow', type=str,required = False)
parser.add_argument('-v', '--system', help='system: j-jimmy s-sarah h-hmtvm w-wwe', type=str) #v for virtual machine

#read in command line arguments
args = parser.parse_args()
model = args.model
cycle = args.cycle

#date/cycle
if args.time != '0':
    date = args.time
    #dta = datetime.utcnow()-timedelta(hours=24*int(args.time))
    dta = datetime.strptime(date,'%Y%m%d')
else:
    dta = datetime.utcnow()
#snow?
if args.snow == '1':
    add_snow = True
else:
    add_snow = False
#what system are we working on?
if args.system == 'j':
    sys = 'jimmy'
elif args.system == 'h':
    sys = 'hmtvm'
elif args.system == 's':
    sys = 'sarah'
elif args.system == 'w':
    sys = 'wwe'
else:
    sys = 'sarah'
###### System Options ######
work, ds, mask_loc = get_system(sys) ###work is where the data is -- ds is where to save things

print(dta,model,args.time,cycle,work)

def fhrs(i,x=2):
    # use zfill method? --with x set as 2 makes a number string with 2 digits
    return str(i).zfill(x)


du = datetime.strftime(dta,'%Y%m%d')
dd = datetime.strftime(dta,'%Y%m%d')+cycle
dj = datetime.strftime(dta,'%y%j')+cycle+'00'

###### Model Options ######
dct,ix,iy,its,ite,wdir = model_dict(model, work, du)
print(dct,ix,iy,its,ite,wdir)

fils = []

if model == 'HREF' or model == 'RRFSCE':
    parms = ['mean','pmmn','lpmm']
    for f, parm in enumerate(parms):
      #file construction
        if model == 'HREF':
            its = 1
            foo = ['href.t'+str(cycle)+'z.conus.'+parm+'.f'+fhrs(i)+'.grib2' for j,i in enumerate(range(its,ite,1))]
            fils.extend(foo)
            print(fils)
        else:
            foo = [model.lower()+'.t'+str(cycle)+'z.conus.'+parm+'.f'+fhrs(i)+'.grib2' for j,i in enumerate(range(its,ite,1))]
            fils.extend(foo)
            print(fils)
elif model == 'SSEF':
    #parms = ['ens', 'sam06', 'sam24']
    parms = ['ens']
    for f, parm in enumerate(parms):
        foo =  ['fv3lam-'+parm+'_'+str(dd)+'f'+fHHH(i)+'.grib2' for j,i in enumerate(range(its,ite,1))]
        fils.extend(foo)
elif model == 'HRRRE' or model == 'HRRREMP':
    fils =  [dj+str(fhrs(i))+'00' for j,i in enumerate(range(its,ite,1))]
else: 
    fils =  [model.lower()+'_ens_'+str(dd)+fhrs(i)+'.grib2' for j,i in enumerate(range(its,ite,1))]

#set mask
mll = np.load(mask_loc+model.lower()+'.npz')
xf = mll['xf']
if xf[0,0] >= 0.:
    xf = xf - 360.
    print('negative longitudes')
yf = mll['yf']
print('SHAPE: ',np.shape(xf),np.shape(yf))
central_lon,central_lat = mll['central']
#central_lon,central_lat = -95, 25

# only precip gets stored in an array for multihour products
precip = np.zeros([ite-its+1,iy,ix])
pprecip = np.zeros([ite-its+1,iy,ix])
lprecip = np.zeros([ite-its+1,iy,ix])
if add_snow:
    snow = np.zeros([ite-its+1,iy,ix])
    catr = np.zeros([ite-its+1,iy,ix])
    catz = np.zeros([ite-its+1,iy,ix])
    cati = np.zeros([ite-its+1,iy,ix])
    cats = np.zeros([ite-its+1,iy,ix])
    if model == 'SSEF':
        psnow = np.zeros([ite-its+1,iy,ix])
        lsnow = np.zeros([ite-its+1,iy,ix])
        homes = np.zeros([1,iy,ix])

goods = []
count = 0
# Read data
for k,fil in enumerate(fils):
#    print("THIS IS K!!!!! ",k)
#    print("THIS IS IK!!!!! ",ik)
#    print("THIS IS COUNT!!!!! ",count)
    good = True
    print(fil)
    if model == 'HREF': # or model == 'RRFSCE':
        if count >= 49: 
            count = 0 
            ik = count
        else:
            ik = count 
    else:
        ik = k 
#    print("THIS IS K!!!!! ",k)
#    print("THIS IS IK!!!!! ",ik)
#    print("THIS IS COUNT!!!!! ",count)
    if 'f000' in fil:
        continue
    else:
        print("not a dumby file")
    try:
        grbs = pygrib.open(dct+fil)
#        print(grbs())
    except:
        print('File not found:',dct+fil)
        good = False
    #extract the initial time and confirm its a match for expected time
    if good:
        try:
            validTime = datetime.strftime(grbs[1].analDate,'%Y%m%d%H')
            if validTime == dd:
                #print(validTime,dd)
                great = 1
            else: #this doesnt do anything duh
                print('Not the correct data file: ',validTime,du)
                raise ValueError 
        except:
            print('Empty grib file')
            continue #was break
        nam = 'Total precipitation'
        grb = [grb for grb in grbs() if grb.parameterName == nam]
        print(grb)
        for grb in grb:
            start = grb.startStep
            end = grb.endStep
            dif = end - start
            print(dif)
            if dif == 1:
                if 'mean' in fil:
                    p = 'precip'
                elif 'pmmn' in fil:
                    p = 'pprecip'
                elif 'lpmm' in fil:
                    p = 'lprecip'
                elif model == 'SSEF':
                    print("THIS IS DIF",dif)
                    try:
                        print("this has a generated process : ",grb.typeOfGeneratingProcess,type(grb.typeOfGeneratingProcess),grb)
                        dfss = (grb.typeOfGeneratingProcess)
                    except:
                        print("SSEF grib does not have a generated process.")
                    if dfss == 4:
                        p = 'precip'
                    elif dfss == 193:
                        p = 'pprecip'
                    elif dfss == 200:
                        p = 'lprecip'
                    else:
                        p = 'prob'
                else:
                    try:
                        print("this is Derived : ",grb.derivedForecast,type(grb.derivedForecast),grb)
                        df = grb.derivedForecast
                    except:
                        print("no dervied value")
                        df = ''
    #                    p = 'prob'
                    if df == 0:
                        p = 'precip'
                    elif df == 1:
                        p = 'pprecip'
                    elif df == 6:
                        p = 'lprecip'
                    else:
                        p = 'prob'
                print(p)
                if p == 'precip' or p == 'pprecip' or p == 'lprecip':
                    if p == 'precip':
                        print("p should be precip --- ",p)
                        precip[ik] = precip[ik-1] + grb.values/25.4
                        print(k,good,'Max: ',np.nanmin(precip[ik]),np.nanmax(precip[ik]))
                        goods.append(good)
                        print(precip[ik],np.shape(precip))
                    elif p == 'pprecip':
                        print("p should be pprecip --- ",p)
                        pprecip[ik] = pprecip[ik-1] + grb.values/25.4
                        print(k,good,'Max: ',np.nanmin(pprecip[ik]),np.nanmax(pprecip[ik]))
                        goods.append(good)
                        print(pprecip[ik],np.shape(pprecip))
                    elif p == 'lprecip':
                        print("p should be lprecip --- ",p)
                        lprecip[ik] = lprecip[ik-1] + grb.values/25.4
                        print(k,good,'Max: ',np.nanmin(lprecip[ik]),np.nanmax(lprecip[ik]))
                        goods.append(good)
                        print(lprecip[ik],np.shape(lprecip))

        if add_snow:
#                winter = ['Water equivalent of accumulated snow depth']
#                winter = ['Water equivalent of accumulated snow depth', 'Categorical rain', 'Categorical freezing rain', 'Categorical ice pellets', 'Categorical snow']
            if model == 'HREF' or model == 'RRFSCE' and 'mean' in fil: # p == 'precip':
                winter = ['Water equivalent of accumulated snow depth', "192","193","194","195"]
                for s,w in enumerate(winter):
                    if w == winter[0]:
                        wname = 'SnDepth'
                        print(wname)
                    elif w == winter[1]:
                        wname = 'crain'
                    elif w == winter[2]:
                        wname = 'cfrzr'
                    elif w == winter[3]:
                        wname = 'cicep'
                    elif w == winter[4]:
                        wname = 'csnow'
                    else:
                        print("not defined")
                    wnam = w
                    print('wnam: ',wnam) 
                    sgrb = [sgrb for sgrb in grbs() if sgrb.parameterName == wnam]
                    #print('SGRB: ',sgrb)
                    for sgrb in sgrb:
                        start = sgrb.startStep
                        end = sgrb.endStep
                        dif = end - start
                        print(dif)
                        if dif == 1:
                            if wname == 'SnDepth':
        #                            snow[ik] = snow[ik-1] + sgrb[indx].values #(sgrb[indx].values/977)*39.3701 ###sgrb[idx].values/997.)*39.3701*10. - is density of water * meter to inches * SLR
                                snow[ik] = snow[ik-1] + ((sgrb.values/997.)*39.3701*10)#*0.0393701*10 ###sgrb[idx].values/997.)*39.3701*10. - is density of water * meter to inches * SLR
                                print(k,good,'Max: ',np.nanmin(snow[ik]),np.nanmax(snow[ik]))
                                goods.append(good)
                                print(snow[ik],np.shape(snow))
                            elif wname == 'crain':
                                catr[ik] = sgrb.values
                            elif wname == 'cfrzr':
                                catz[ik] = sgrb.values
                            elif wname == 'cicep':
                                cati[ik] = sgrb.values
                            elif wname == 'csnow':
                                cats[ik] = sgrb.values
                            else:
                                print("not defined")
            elif model == 'SSEF':
                wnam = 'Total snowfall'
                wname = 'TotSnow'
                sgrb = [sgrb for sgrb in grbs() if sgrb.parameterName == wnam]
                for sgrb in sgrb:
                    start = sgrb.startStep
                    end = sgrb.endStep
                    dif = end - start
                    if dif == 1:
                        print("this is k: ",k)
                        print("this is ik: ",ik)
                        try:
                            print("this has a generated process : ",sgrb.typeOfGeneratingProcess,type(sgrb.typeOfGeneratingProcess),grb)
                            dfss = (sgrb.typeOfGeneratingProcess)
                        except:
                            print("SSEF grib does not have a generated process.")
                            dfss = ''
                        if dfss == 4:
                           p = 'snow'
                        elif dfss == 193:
                            p = 'psnow'
                        elif dfss == 200:
                           p = 'lsnow'
                        else:
                           p = 'prob'
                        if p == 'snow':
                            homes = (sgrb.values*100)/2.54
#                            print("HERE - first homes ",np.unique(homes))
                            homes[homes==9999] = 0.0 
#                            print("HERE - secon homes ",np.unique(homes))
                            homes[np.isnan(homes)] = 0
#                            print("HERE - third homes ",np.unique(homes))
                            snow[ik] = snow[ik-1] + homes # (sgrb.values*100)/2.54 #SSEF in meters
                            print("this is p", p)
                        elif p == 'psnow':
                            homes = (sgrb.values*100)/2.54
                            homes[homes==9999] = 0.0 
                            homes[np.isnan(homes)] = 0
                            psnow[ik] = psnow[ik-1] + homes #  (sgrb.values*100)/2.54 
                            print("this is p", p)
                        elif p == 'lsnow':
                            homes = (sgrb.values*100)/2.54
                            homes[homes==9999] = 0.0 
                            homes[np.isnan(homes)] = 0 
                            lsnow[ik] = lsnow[ik-1] + homes # (sgrb.values*100)/2.54 
                            print("this is p", p)

    count += 1


if add_snow:
    np.savez_compressed(dct+model.lower()+'_mean_'+dd+'.npz',xf=xf,yf=yf,fcsts=precip,snow=snow,central = np.array([central_lon,central_lat]))  
    if model == 'SSEF':
        np.savez_compressed(dct+model.lower()+'_pmm_'+dd+'.npz',xf=xf,yf=yf,fcsts=pprecip,snow=psnow,central = np.array([central_lon,central_lat]))
        np.savez_compressed(dct+model.lower()+'_lpmm_'+dd+'.npz',xf=xf,yf=yf,fcsts=pprecip,snow=lsnow,central = np.array([central_lon,central_lat]))
    elif model == 'HREF' or model == 'RRFSCE':
        np.savez_compressed(dct+model.lower()+'_categorical_'+dd+'.npz',xf=xf,yf=yf,crain=catr,cfrzr=catz,cicep=cati,csnow=cats,central = np.array([central_lon,central_lat]))  
        np.savez_compressed(dct+model.lower()+'_pmm_'+dd+'.npz',xf=xf,yf=yf,fcsts=pprecip,central = np.array([central_lon,central_lat]))
        np.savez_compressed(dct+model.lower()+'_lpmm_'+dd+'.npz',xf=xf,yf=yf,fcsts=lprecip,central = np.array([central_lon,central_lat]))
    else:
        np.savez_compressed(dct+model.lower()+'_pmm_'+dd+'.npz',xf=xf,yf=yf,fcsts=pprecip,central = np.array([central_lon,central_lat]))
        np.savez_compressed(dct+model.lower()+'_lpmm_'+dd+'.npz',xf=xf,yf=yf,fcsts=lprecip,central = np.array([central_lon,central_lat]))
else:
    np.savez_compressed(dct+model.lower()+'_mean_'+dd+'.npz',xf=xf,yf=yf,fcsts=precip,central = np.array([central_lon,central_lat]))
    np.savez_compressed(dct+model.lower()+'_pmm_'+dd+'.npz',xf=xf,yf=yf,fcsts=pprecip,central = np.array([central_lon,central_lat]))
    np.savez_compressed(dct+model.lower()+'_lpmm_'+dd+'.npz',xf=xf,yf=yf,fcsts=lprecip,central = np.array([central_lon,central_lat]))

#np.savez_compressed(work+model.lower()+'.npz',xf=xf,yf=yf,fcsts=lprecip,central = np.array([central_lon,central_lat]))