from netCDF4 import Dataset
import netCDF4 as nc
import numpy as np
import xarray as xr
from wrf import getvar, vinterp, ALL_TIMES
import os
from datetime import datetime
import dask
from dask.distributed import Client

client = Client(n_workers=16)
@dask.delayed
def WriteSurfVariables(inputDir,outDir,file):
    inputnc=Dataset(inputDir+"/"+file)
    varlist=['slp','T2','td2','uvmet10','Q2','PSFC','RAINC','RAINNC','SWDOWN','GLW','HFX',
        'LH','OLR','SNOW','SST','QFX']
    for var in varlist:
        outvar=getvar(inputnc,var,timeidx=ALL_TIMES)
        del outvar.attrs['projection']
        outvar.to_netcdf(outDir+var+"_"+file) 
#n_worker=16
now = datetime.now() 
print(now) 
#results2d = []
results3d = []
inputDir="/blue/dhingmire/Prod_Outputs/ERA5/1981/d2"
outDir="/blue/dhingmire/Processed_data/ERA5/testDir/"
for file in os.listdir(inputDir):
    
    results3d.append(WriteSurfVariables(inputDir,outDir,file))
    #results3d.append(Write3DVariables(inputDir,outDir,file))

#dask.compute(results2d)
dask.compute(results3d)

    
now = datetime.now()
print("now =", now)