from netCDF4 import Dataset
import numpy as np
import xarray as xr
import os
import glob

# Ensemble member(s) to be processed
mem_arr = ["rtty_1990"]
numMEM = len(mem_arr)

# Years we are interested in
years_arr = ["1990","2003"]
numYEARS = len(years_arr)

# Loop through the number of ensemble members
for mem in mem_arr:

# Loop through the number of years
    for years in years_arr:

        print("Year = {}".format(years))
# Create the netdcf output file
        outname = "t2_maxmin_daily_jja_"+years+"_"+mem+".nc"
        os.system("rm -rf t2_maxmin_daily_jja_"+years+"_"+mem+".nc")
        out = Dataset("t2_maxmin_daily_jja_"+years+"_"+mem+".nc","w")

# Find the data files for the year of interest, list them,
#   and put them in an array of filenames to be read
        DATADir = "/sysdisk1/C3WE/tutorial/model_data/"
        files_arr = sorted(glob.glob(DATADir+"auxhist1_d01_"+years+"*"))
        numFILES = len(files_arr)

        tmaxday = np.full((numFILES,279,395),9.96921e+36)
        tlandmaxday = np.copy(tmaxday)
        tminday = np.copy(tmaxday)
        tlandminday = np.copy(tmaxday)

# Because wrf is sometime weird, there is a problem with the landmask in the output
#   files.  Therefore we are taking the landmask from a met_em wrf input file.
        a = Dataset("data/met_em.d01.2004-12-16_12:00:00.nc","r")
        landmask = a.variables['LANDMASK'][0,:,:]

# Looping through a years worth of files
        day=0
        for files in files_arr:

# Reading in a day at a time and moving the file attributes to our output file
            a = Dataset(files,"r")
            if day == 0:
                glob_atts = a.__dict__
                out.setncatts(glob_atts)

# Read in temperature, latitude and longitude
            tk = a.variables['T2']
            t = tk[:]-273.15
            lat = a.variables['XLAT'][0,:,:]
            lon = a.variables['XLONG'][0,:,:]

# Masking out the oceans using the landmask to just get t2 over land
            tland = np.copy(t)
            tland[:,landmask == 0] = np.nan

# Calculated maxes for t2 and t2mask
            tmax = np.max(t,axis=0)
            tlandmax = np.nanmax(tland,axis=0)

# Writing the maxes to the maxday variables
            tmaxday[day,:,:] = tmax
            tlandmaxday[day,:,:] = tlandmax

# Doing the same for minimums...
            tmin = np.min(t,axis=0)
            tlandmin = np.nanmin(tland,axis=0)

            tminday[day,:,:] = tmin
            tlandminday[day,:,:] = tlandmin
            day=day+1

# Writing out lat, lon, and landmask to the netcdf file
        out.createDimension('Time',None)
        out.createDimension('south_north',len(a.dimensions['south_north']))
        out.createDimension('west_east',len(a.dimensions['west_east']))

        XLAT = out.createVariable('XLAT',np.dtype(tk),['south_north','west_east'])
        XLAT.setncatts({'FieldType': u'104','MemoryOrder': u'XY',\
                        'description': u'LATITUDE, SOUTH IS NEGATIVE','units': u'degree_north'})

        XLONG = out.createVariable('XLONG',np.dtype(tk),['south_north','west_east'])
        XLONG.setncatts({'FieldType': u'104','MemoryOrder': u'XY',\
                        'description': u'LONGITUDE, WEST IS NEGATIVE','units': u'degree_east'})

        LANDMASK = out.createVariable('LANDMASK',np.dtype(tk),['south_north','west_east'])
        LANDMASK.setncatts({'FieldType': u'104','MemoryOrder': u'XY',\
                        'description': u'Landmask : 1=land, 0=water','units': u'none'})

        out.variables['XLAT'][:]=lat[:]
        out.variables['XLONG'][:]=lon[:]
        out.variables['LANDMASK'][:]=landmask[:]

# Writing out the variable attributes and values to the netcdf output file

        T2MAX = out.createVariable('T2MAX',np.dtype(tk),('Time','south_north','west_east'))
        T2MAX.setncatts({'description': u'DAILY MAX TEMP at 2M','coordinates': u'XLONG XLAT',\
                         'units': u'C','FillValue': u'9.96921e+36'})

        T2MAX_MASK = out.createVariable('T2MAX_MASK',np.dtype(tk),('Time','south_north','west_east'))
        T2MAX_MASK.setncatts({'description': u'DAILY MAX TEMP at 2M (masked)','coordinates': u'XLONG XLAT',\
                              'units': u'C','FillValue': u'9.96921e+36'})

        T2MIN = out.createVariable('T2MIN',np.dtype(tk),('Time','south_north','west_east'))
        T2MIN.setncatts({'description': u'DAILY MIN TEMP at 2M','coordinates': u'XLONG XLAT',\
                         'units': u'C','FillValue': u'9.96921e+36'})

        T2MIN_MASK = out.createVariable('T2MIN_MASK',np.dtype(tk),('Time','south_north','west_east'))
        T2MIN_MASK.setncatts({'description': u'DAILY MIN TEMP at 2M (masked)','coordinates': u'XLONG XLAT',\
                              'units': u'C','FillValue': u'9.96921e+36'})

        out.variables['T2MAX'][:]=tmaxday[:]
        out.variables['T2MAX_MASK'][:]=tlandmaxday[:]
        out.variables['T2MIN'][:]=tminday[:]
        out.variables['T2MIN_MASK'][:]=tlandminday[:]

        out.close()