; Subsetting PRISM files to Great Plains states and calculate state average for NIFA project

; ---------------------------------------------------

load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_csm.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/contributed.ncl"

; ---------------------------------------------------

begin

state_list = [/"west","east"/]

state_title = [/"West of the 100th Meridian","East of the 100th Meridian"/]

variable_list = [/"ppt","tavg"/]
; variable_list = [/"ppt","tavg","tmax","tmin","vpr"/]

units_list = [/"mm per year","degrees C"/]
; units_list = [/"mm per year","degrees C","degrees C","degrees C","Pa"/]

long_name_list = [/"Annual total precipitation","Annual average mean air temperature"/]
; long_name_list = [/"Annual total precipitation","Annual average mean air temperature","Annual average maximum air temperature","Annual average minimum air temperature","Annual average vapor pressure deficit"/]

lat_max = [/49.1,49.5/]
lat_min = [/27.7,25.7/]
lon_max = [/-99.5,-89.3/]
lon_min = [/-116.2,-100.5/]

do v = 0,ListCount(variable_list) - 1

  ; ***********************************************
  ; 1 - READ IN PRISM DATA AND SHAPEFILE
  ; ***********************************************

  print(systemfunc("date"))

  print("1 - Read in file ~/NIFA/MACAv2/PRISM/" + variable_list[v](0) + "_1951_2014.nc")

	PRISM_f = addfile("~/NIFA/MACAv2/PRISM/" + variable_list[v](0) + "_1951_2014.nc","r")


  ; ***********************************************
  ; 2 - READ IN VARIABLES
  ; ***********************************************

  print("2 - Read in variable for 1970-2012")

  if(variable_list[v](0) .eq. "ppt") then
;  	variable = PRISM_f->ppt(228:743,:,:)
  	variable = PRISM_f->ppt(228:239,:,:)
  end if

  if(variable_list[v](0) .eq. "tavg") then
  	variable = PRISM_f->tas(228:743,:,:)
  end if

  if(variable_list[v](0) .eq. "tmax") then
  	variable = PRISM_f->tmax(228:743,:,:)
  end if

  if(variable_list[v](0) .eq. "tmin") then
  	variable = PRISM_f->tmin(228:743,:,:)
  end if

  if(variable_list[v](0) .eq. "vpr") then
  	variable = PRISM_f->vpr(228:743,:,:)
  end if

  lat = PRISM_f->lat
  lon = PRISM_f->lon
  description = long_name_list[v](0)
  units = units_list[v](0)

;	time = ispan(1970,2012,1)							; list with years in question (1970-2012)
	time = ispan(1970,1970,1)							; list with years in question (1970-2012)

  time!0 = "time"
  time&time = time
	time@long_name = "calendar year"
  time@units = "year"


  ; ***********************************************
  ; 3 - CREATE STATE SUBSET OF THE GLOBAL DATA
  ; ***********************************************

	print("3 - Create state lat/lon subset")

do s = 0,ListCount(state_list) - 1

  print("3a - Read in shapefile for  ---  " + state_title[s](0) + "  ---")

	shp = addfile("~/NIFA/shapefile/states/8x_" + state_list[s](0) + ".shp", "r")

  ; let's choose latitudes of 
    latind = (/lat_min[s],lat_max[s]/)						; latind is a variable

  ; let's choose longitudes of 
    lonind = (/lon_min[s],lon_max[s]/)						; lonind is a variable

  lti = ind_nearest_coord(latind,lat,0)				; lti,lni = variables, ind_nearest_coord = pre-defined functions to determine the indices of locations closest to the coordinate array, needed later to define var_annual_3D
  lni = ind_nearest_coord(lonind,lon,0)

  lat_subdomain = lat(lti(0):lti(1))
  lon_subdomain = lon(lni(0):lni(1))
  
;  printMinMax(decimalPlaces(lat_subdomain,2,True),False)
;  printMinMax(decimalPlaces(lon_subdomain,2,True),False)

  var_lat_lon_3D = variable(:,lti(0):lti(1),lni(0):lni(1))	; define the subdomain based on latind and lonind


  ; *************************************************
  ; 4 - CREATE ANNUAL TOTAL/AVERAGE FROM MONTHLY DATA
  ; *************************************************

  print("4 - Create annual total/average from monthly data")
  
  var_lat_lon_3D_annual = new((/43,dimsizes(lat_subdomain),dimsizes(lon_subdomain)/),float,variable@_FillValue)
  
  m = 0
  
;	do y = 0,42
	do y = 0,0
		
		if (variable_list[v](0) .eq. "ppt") then					; Sum for precip, average for all other variables, to get annual data from monthly data
			var_lat_lon_3D_annual(y,:,:) = dim_sum_n(var_lat_lon_3D(m:m+11,:,:),0)
		else
			var_lat_lon_3D_annual(y,:,:) = dim_avg_n(var_lat_lon_3D(m:m+11,:,:),0)
		end if

		m = m + 12
		
	end do
	
	delete(y)
	
;	printVarSummary(var_lat_lon_3D_annual)
;	printMinMax(var_lat_lon_3D_annual,False)


  ; ***********************************************
  ; 5 - CREATE SUBSET WITH SHAPEFILE INFORMATION
  ; *********************************************** 

  print("5 - Create subset based on shapefile")

	shp_lat = decimalPlaces(shp->y,5,True)
  shp_lon = decimalPlaces(shp->x,5,True)

  shp_lat!0 = "shp_lat"
  shp_lat&shp_lat = shp_lat
  shp_lat@long_name = "latitude"
  shp_lat@units = "degrees_north"

  shp_lon!0 = "shp_lon"
  shp_lon&shp_lon = shp_lon
  shp_lon@long_name = "longitude"
  shp_lon@units = "degrees_east"

  nshp = dimsizes(shp_lon)
  
; ***** Define lat/lon extent of the state in question
  min_shp_lat = min(shp_lat)
  max_shp_lat = max(shp_lat)
  min_shp_lon = min(shp_lon)
  max_shp_lon = max(shp_lon)

; ***** Determine the boundary lat/lons of the subset region within the lat/lon box
  ilt_min = ind(min_shp_lat.gt.lat_subdomain)
  ilt_max = ind(max_shp_lat.lt.lat_subdomain)
  iln_min = ind(min_shp_lon.gt.lon_subdomain)
  iln_max = ind(max_shp_lon.lt.lon_subdomain)
  ilt1   = ilt_min(dimsizes(ilt_min)-1)    		; Start of lat box
  ilt2   = ilt_max(0)                     		; End of lat box
  iln1   = iln_min(dimsizes(iln_min)-1)    		; Start of lon box
  iln2   = iln_max(0)                     		; End of lon box
  
;	print((/ilt1/))
;	print((/ilt2/))
;	print((/iln1/))
;	print((/iln2/))
;	sleep(10)

; ***** Create new subset based on shapefile ***** (source: https://www.ncl.ucar.edu/Applications/Scripts/mask_9.ncl)

  MASK_INSIDE = True    ; Whether to mask ("protect") data inside (True) or outside (False) the given geographical area

  if (MASK_INSIDE) then						; start with all missing data -> then fill in data inside the shapefile

		print("     - MASK_INSIDE -")

		var_shp_3D_annual = new(dimsizes(var_lat_lon_3D_annual),typeof(var_lat_lon_3D_annual),var_lat_lon_3D_annual@_FillValue)
		
;		printVarSummary(var_shp_3D_annual)
;		printMinMax(var_shp_3D_annual,False)
		
		copy_VarCoords(var_lat_lon_3D_annual,var_shp_3D_annual)
;		do y = 0,42
		do y = 0,0
			print((/time(y)/))
			do ilt=ilt1,ilt2
				do iln=iln1,iln2
					if(gc_inout(lat_subdomain(ilt),lon_subdomain(iln),shp_lat,shp_lon)) then	; returns TRUE if lat/lon point is *inside* the polygon
						var_shp_3D_annual(y,ilt,iln) = var_lat_lon_3D_annual(y,ilt,iln)         ; replace missing values *inside* the study region with data again
        	end if
    	  end do
  	  end do
  	end do
	end if
 	
; 	printMinMax(var_shp_3D_annual,False)


; ***** Define lat/lon of the state in question

	latind_shp = (/min_shp_lat,max_shp_lat/)
	lonind_shp = (/min_shp_lon,max_shp_lon/)

	lti_shp = ind_nearest_coord(latind_shp,lat,0)				; lti,lni = variables, ind_nearest_coord = pre-defined functions to determine the indices of locations closest to the coordinate array, needed later to define var_annual_3D
	lni_shp = ind_nearest_coord(lonind_shp,lon,0)

  lat_state = lat(lti_shp(0):lti_shp(1))
  lon_state = lon(lni_shp(0):lni_shp(1))

  var_state = var_shp_3D_annual(:,:,:)	; define the subdomain based on latind and lonind
;  var_state = var_shp_3D_annual(:,lti_shp(0):lti_shp(1),lni_shp(0):lni_shp(1))	; define the subdomain based on latind and lonind
	
;	printVarSummary(var_state)
;	printMinMax(var_state,False)


  ; ***********************************************
  ; 6 - CALCULATE ANNUAL TOTAL/AVERAGE
  ; ***********************************************

  print("6 - Calculate annual total/average")

	annual_avg_help = dim_avg_n(var_state,2)
	annual_avg = dim_avg_n(annual_avg_help,1)

	annual_avg@long_name = description
	annual_avg@units = units_list[v](0)
	
;	printVarSummary(annual_avg)
;	printMinMax(annual_avg,False)


  ; ***********************************************
  ; 7 - WRITE THE SUBSET INTO NEW netCDF FILE
  ; ***********************************************

  print("7 - Write subset into new netCDF file")

;  printVarSummary(var_state)
;  printVarSummary(lat_subdomain)
;  printVarSummary(lon_subdomain)
;  printVarSummary(time)

    netCDF1 = True						; Output format is NetCDF

  if (netCDF1) then
    diro = "~/NIFA/MACAv2/subset_shp/PRISM_E-W_annually/"
    filo = variable_list[v](0) + "_shp_8x_100_" + state_list[s](0) + ".nc"	; output filename for subset using initialization year and month to differentiate

  end if

  setfileoption("nc","Format","LargeFile")

  if (netCDF1) then
    system("/bin/rm -f "+ diro + filo)				; removes old files with the same name if they are present
    ncdf = addfile(diro + filo,"c")

    setfileoption(ncdf,"DefineMode",True)

    ; create attributes
    fAtt = True
;    fAtt@title = "Data: PRISM data 1970-2012 with annual total (ppt) or average (tavg, tmin, tmax, vpr), Variable: " + variable_list[v](0) + ", Region: " + state_title[s](0)
    fAtt@Conventions = "None"
    fAtt@creation_date = systemfunc("date")

    fileattdef(ncdf,fAtt)

    dimNames = (/"time","lat","lon"/)
    dimSizes = (/dimsizes(time),dimsizes(lat_subdomain),dimsizes(lon_subdomain)/)
    dimUnlim = (/True,False,False/)

    filedimdef(ncdf,dimNames,dimSizes,dimUnlim)

    filevardef(ncdf,"time",typeof(time),(/"time"/))
    filevardef(ncdf,"lat",typeof(lat_subdomain),(/"lat"/))
    filevardef(ncdf,"lon",typeof(lon_subdomain),(/"lon"/))
    filevardef(ncdf,"var_state",typeof(var_state),(/"time","lat","lon"/))

    ; now write all the variables to the file

		ncdf->time 					= (time)
    ncdf->lat_subdomain = (lat_subdomain)
    ncdf->lon_subdomain	= (lon_subdomain)
    ncdf->var_state	 		= (var_state)
    ncdf->annual_avg 		= (annual_avg)

  end if

  delete([/var_lat_lon_3D_annual,var_lat_lon_3D,var_shp_3D_annual,lat_subdomain,lon_subdomain/])
  
  delete([/shp_lat,shp_lon,nshp,min_shp_lat,max_shp_lat,min_shp_lon,max_shp_lon,ilt_min,ilt_max,iln_min,iln_max,ilt1,ilt2,iln1,iln2/])
  
  delete([/latind_shp,lonind_shp,lti_shp,lni_shp,lat_state,lon_state,var_state/])
  
  delete([/annual_avg_help,annual_avg/])

  delete([/filo,diro,ncdf,fAtt/])

;  print("*** TEST DONE *** PAUSE ***")
  
  print(systemfunc("date"))
  
;  sleep(10000)

end do	; end do loop state shapefile

	delete([/variable,lat,lon,time,description,units/])

end do	; end do loop variable

  print("**** ALL DONE ****")

end