##On my laptop I installed PyNGL and PyNIO to use Nio and Ngl module for reading *.nc file and for contour plot
##On titan, I did not install PyNGL and PyNIO, so import Dataset from netCDF4 to read *.nc file
##probably, I will the following command for contourplot in titan python "from mpl_toolkits.basemap import Basemap"


#!/usr/bin/env python3

import numpy as np
import numpy.ma as ma
import os
from netCDF4 import Dataset


# ---------------- Data sources ------------------------
input_file1 = '/lustre/atlas1/cli115/world-shared/4ue/20180215.DECKv1b_H1.ne30_oEC.edison/postproc/lnd/climo_regrid/20180215.DECKv1b_H1.ne30_oEC.edison_JJA_climo.nc'
input_file3 = '/lustre/atlas1/cli115/world-shared/4ue/20180215.DECKv1b_H1.ne30_oEC.edison/postproc/lnd/climo_regrid/20180215.DECKv1b_H1.ne30_oEC.edison_ANN_climo.nc'

# --------------------------------------------------------------
# --------------------------------------------------------------

# input_file1 get FLDS (mean downwelling longwave rad),FIRA (mean net IR longwave rad, absorbed)
# FLDS(0,:,:), FIRA(0,:,:)
# Matrix gris_area(Greenland ice area) has values 0 and none

# Goal: print out area-weighted average of SW_d(268), SW_net(61), LW_d(235), LW_net(-46), R_net(15), SHF(7), LHF(-8),
# these variables have size[192,288]
# FSDS: atmospheric incident solar radiation
# FSA: absorbed solar radiation
# FLDS: atmospheric longwave radiation
# FIRA: net infrared (longwave) radiation
# FIRE: emitted infrared (longwave) radiation
# FSH: sensible heat
# QSOIL: ground evaporation over ice landunits


# ---read input_file2, probably get area with ice sheet with nonzero values
#ncid2 = Dataset(input_file2, mode='r')
#ices = ncid2.variables['PCT_GLC_ICESHEET'][:] 
#ncid2.close()

# ---read input_files
ncid1 = Dataset(input_file1, mode='r')
#ncid11 = Dataset(input_file11, mode='r')
ncid3 = Dataset(input_file3, mode='r')

FSDS = ncid1.variables['FSDS'][0,:,:] 
FSA = ncid1.variables['FSA'][0,:,:] 
FLDS = ncid1.variables['FLDS'][0,:,:] 
FIRA = ncid1.variables['FIRA'][0,:,:]
FSH = ncid1.variables['FSH'][0,:,:] 
QSOIL = ncid1.variables['QSOIL'][0,:,:] 
RAIN = ncid3.variables['RAIN'][0,:,:] 
SNOW = ncid3.variables['SNOW'][0,:,:] 
QRUNOFF = ncid3.variables['QRUNOFF'][0,:,:] 
#T2m = ncid11.variables['TREFHT'][0,:,:] 
T2m = ncid1.variables['TSA'][0,:,:] 

#E3SM
GrisMask = ncid1.variables['gris_mask'][:,:] 

SW_d  = FSDS 
SW_net = FSA
LW_d = FLDS
LW_net = -1.0 * FIRA
R_net = FSA-FIRA
SHF = -1.0 * FSH
LHF = -1.0 * 2.835e6 * QSOIL
#SMB = QICE * 3.1567e7
SMB = RAIN + SNOW - QRUNOFF - QSOIL * 3.1567e7
T2mC = T2m - 273.15

# ---grid cell area on Greenland. For other regions, the value is 9.999e35.
area = ncid1.variables['area'][:] 

ncid1.close()
#ncid11.close()
ncid3.close()

# # # --- Option I: use greenland mask to mask the area where the ice sheet fraction is zero
# # # Note: some non-greenland region has non-zero ice sheet fraction.
Gris_mask = ma.masked_equal(GrisMask,0)
area_maskice = ma.masked_array(area, mask=Gris_mask.mask)


# # # --- Option II: use icesheet fraction mask to mask the area where the ice sheet fraction is zero
# ices_mask = ma.masked_equal(ices,0)
# area_maskice = ma.masked_array(area, mask=ices_mask.mask)


# # # --- Option III: first multiply area with ices, then use greenland mask to mask the area where the ice sheet fraction is zero
## -- this option gives the closest results to the Table 1 in the paper.
#areaices = np.multiply(area,ices)
#Gris_mask = ma.masked_equal(GrisMask,0)
#area_maskice = ma.masked_array(areaices, mask=Gris_mask.mask)


swdout=(ma.average(SW_d,weights=area_maskice))
swnetout=(ma.average(SW_net,weights=area_maskice))
lwdout=(ma.average(LW_d,weights=area_maskice))
lwnetout=(ma.average(LW_net,weights=area_maskice))
rnetout=(ma.average(R_net,weights=area_maskice))
shfout=(ma.average(SHF,weights=area_maskice))
lhfout=(ma.average(LHF,weights=area_maskice))
smbout=(ma.average(SMB,weights=area_maskice))
t2cout=(ma.average(T2mC,weights=area_maskice))

print "Sw_d model:", swdout
print "Sw_net model:", swnetout
print "Lw_d model:", lwdout
print "Lw_net model:", lwnetout
print "Rnet model:", rnetout
print "SHF model:", shfout
print "LHF model:", lhfout
print "SMB model:", smbout
#print "T2M C atm model:", t2cout
print "T2M C land model:", t2cout