###########################################################################################
# Script to compute SPEI (Thornthwaite)rasters using PRISM 800m
# Using precipitation and temperature from PRISM800m rasters to calculate
# Thornthwaite-based SPEI
# @USDA-ARS-Tucson, AZ
# @Date: 04-20-2016
###########################################################################################

library(raster)
library(data.table)

### PRECIPITATION DATA
setwd('~/smbdir/PRISM800m/MONTHLY/ppt/')
# Get a vector with all daiy files
all.files <- list.files(pattern='\\.tif$')
# Get a stack and convert it into a matrix -> data.table
ptm<-proc.time()
ppt.stack <- stack(all.files)
mt <- as.matrix(ppt.stack)
DT <- as.data.table(mt)
proc.time() - ptm
print("Done...")

# Get x,y from converting raster template to points

rTemp <- raster('cai_ppt_us_us_30s_201312.tif')
rTemp[1:ncell(rTemp)] <- 1

rTemp <- rasterToPoints(rTemp)
rTemp <- as.data.table(rTemp)

# Add LATITUDE and LONGITUDE columns to DT
DT[,LATITUDE:=rTemp$y]
DT[,LONGITUDE:=rTemp$x]
dt.names <- c(names(DT)[1369:1370],names(DT)[1:1368])
setcolorder(DT,dt.names)

# Add PIXELID column and re-arrange columns
DT[,PIXELID:=seq(1:nrow(DT))]
dt.names <-c(names(DT)[nrow(DT)], names(DT)[1:1371])
setcolorder(DT,dt.names)

# Change names of dates
setnames(DT,names(DT)[4:ncol(DT)],paste0("D_",substr(names(DT)[4:ncol(DT)],19,22),substr(names(DT)[4:ncol(DT)],23,24),"01"))

saveRDS(DT,'~/DATA/DATA_TABLES/dt.prism.ppt.1900.2013.rds')

### TEMPERATURE DATA##########################
#############################################

setwd('~/smbdir/PRISM800m/MONTHLY/tmean/')
# Get a vector with all daiy files
all.files <- list.files(pattern='\\.tif$')
# Get a stack and convert it into a matrix -> data.table
ptm<-proc.time()
temp.stack <- stack(all.files)
mt <- as.matrix(temp.stack)
DT <- as.data.table(mt)
proc.time() - ptm
print("Done...")

# Get x,y from converting raster template to points

rTemp <- raster('cai_tmean_us_us_30s_201312.tif')
rTemp[1:ncell(rTemp)] <- 1

rTemp <- rasterToPoints(rTemp)
rTemp <- as.data.table(rTemp)

# Add LATITUDE and LONGITUDE columns to DT
DT[,LATITUDE:=rTemp$y]
DT[,LONGITUDE:=rTemp$x]
dt.names <- c(names(DT)[1369:1370],names(DT)[1:1368])
setcolorder(DT,dt.names)

# Add PIXELID column and re-arrange columns
DT[,PIXELID:=seq(1:nrow(DT))]
dt.names <-c(names(DT)[ncol(DT)], names(DT)[1:1370])
setcolorder(DT,dt.names)

# Change names of dates
setnames(DT,names(DT)[4:ncol(DT)],paste0("D_",substr(names(DT)[4:ncol(DT)],21,24),substr(names(DT)[4:ncol(DT)],25,26),"01"))

saveRDS(DT,'~/DATA/DATA_TABLES/dt.tmean.1900.2013.rds')


DT <- readRDS('~/DATA/DATA_TABLES/dt.tmean.1900.2013.rds')

### GET PET ###
setkey(DT,PIXELID)
DT.temp <- DT[!is.na(D_19000101),]
# Set water year from oct-sep
#DT.temp[,c(4:12):=NULL, with=FALSE]
#DT.temp[,c(1360:1362):=NULL, with=FALSE]
# Create ref. template
DT.ref <- DT[,list(PIXELID,LATITUDE,LONGITUDE)]
setkey(DT.ref, PIXELID)
rm(DT)

# GET numeric vector to use with ts
vcols <- names(DT.temp)[c(2,4:ncol(DT.temp))]

# Join monthly data into a column-list, each cell contains the monthly time series for each pixel
ptm<-proc.time()
DT.temp[,TS_COL:=list(list(list(as.numeric(.SD)))),.SDcols=vcols,by=PIXELID]  # 1871.694 seconds
proc.time() - ptm
DT.temp <- DT.temp[,list(PIXELID,LATITUDE,TS_COL)]
gc()
#
#saveRDS(DT.temp,'~/DATA/DATA_TABLES/dt.temp.TS_COL.rds')
#DT.temp <- readRDS('~/DATA/DATA_TABLES/dt.temp.TS_COL.rds')

# Parallel setup
library(parallel)

vCluster <- makeCluster(65)

GetTempTS <- function(x) {
	l <- unlist(x)
	list.final <- list(lat=l[1], tst=ts(l[2:length(l)], frequency=12, start=c(1900,1)))
}
GetPET <- function(x) {
	thornthwaite(x$tst, x$lat, na.rm=TRUE)
}

GetDTpet <- function (x) {
	#transpose(setDT(list(x)))
	as.numeric(t(as.matrix(unlist(x))))
}

GetWB <- function(x) {
       list(wb=ts(unlist(x), frequency=12, start=c(1900,1)))
}
GetSPEI <- function(x) {
	spei(x$wb, 12)$fitted
}

# Make libraries and functions available to cluster
clusterEvalQ(vCluster, c(library(data.table), library(SPEI)))
clusterExport(vCluster, c('GetTempTS', 'GetPET', 'GetDTpet', 'GetSPEI','GetWB'))

# Prepare Timeseries of temperature to use as input for thornthwaite
ptm<-proc.time()
list.ts <- parLapply(vCluster, DT.temp$TS_COL, GetTempTS)  # 3183 seconds
proc.time() - ptm


stopCluster(vCluster)
rm(vCluster)
rm(DT.temp)


# Get PET
ptm<-proc.time()
list.pet <- parLapply(vCluster, list.ts, GetPET)  # 4787.603 seconds
proc.time() - ptm

#saveRDS(list.pet, '~/smbdir/DATA_TABLES/list.pet.rds')  # This is 100 GB file, taking quite of time saving.
ptm <- proc.time()
list.pet2 <- readRDS('~/smbdir/DATA_TABLES/list.pet.rds')  # 1087 seconds
proc.time() - ptm
#
ptm<-proc.time()
list.pet.dt <- parLapply(vCluster, list.pet, GetDTpet) # 2450 seconds with 68 cores
proc.time() - ptm
#
# GetDTpet -  Convert list of vectors with PET into data.table
# Maybe it's a good idea to stop the cluster and remove list.pet to release memory
stopCluster(vCluster)
rm(list.pet)
gc()
ptm<-proc.time()
dt.pet <-as.data.table(t(as.data.table(list.pet.dt)))  # 3482 seconds
proc.time() - ptm
#

# Open precipitation data to substract dt.pet
dt.ppt <- readRDS('~/DATA/DATA_TABLES/dt.ppt.1900.2013.rds')
# Remove NA's
dt.ppt <- dt.ppt[!is.na(D_19000101),]
dt.ppt.cols <- dt.ppt[,4:ncol(dt.ppt),with=FALSE]
# Get waterbalance (PPT - PET)
dt.wb <- dt.ppt.cols - dt.pet
dt.wb[,c('PIXELID','LATITUDE','LONGITUDE'):=dt.ppt[,list(PIXELID,LATITUDE,LONGITUDE)]]

vcols <- names(dt.wb)[1:1368]
vcols <- c('PIXELID','LATITUDE','LONGITUDE',vcols)
setcolorder(dt.wb,vcols)
setkey(dt.wb, PIXELID)

saveRDS('~/DATA/DATA_TABLES/dt.wb.prism.1900.2013.rds')

## set WB time series
vcols <- names(dt.wb)[c(1,4:ncol(dt.wb))]
ptm<-proc.time()
dt.wb[,TS_COL:=list(list(list(as.numeric(.SD)))),.SDcols=vcols,by=PIXELID]  # 2806 seconds
proc.time() - ptm
dt.wb <- dt.wb[,list(PIXELID,TS_COL)]

# Setup Water balance time series
ptm<-proc.time()
list.wb.ts <- parLapply(vCluster, dt.wb$TS_COL, GetWB)  # 3183 seconds
proc.time() - ptm

# Get SPEI time series
ptm<-proc.time()
list.spei.ts <- parLapply(vCluster, list.wb.ts, GetSPEI)  # 11613 seconds (3.2 hours)
proc.time() - ptm
#
stopCluster(vCluster)
rm(vCluster)
gc()

vCluster <- makeCluster(65)
# Get a data table with SPEI
ptm<-proc.time()
list.spei.dt <- parLapply(vCluster, list.spei.ts, GetDTpet) #1116.363
proc.time() - ptm
#
ptm<-proc.time()
dt.spei <-as.data.table(t(as.data.table(list.spei.dt)))  # 3482 seconds
proc.time() - ptm
#
ptm<-proc.time()
saveRDS(dt.spei,'~/DATA/DATA_TABLES/dt.spei.conus.1901_2013.rds')
proc.time() - ptm
# Saving time: 5.9 hrs

# Get PIXELID, LAT,LNG
DT <- readRDS('~/DATA/DATA_TABLES/dt.tmean.1900.2013.rds')[,list(PIXELID,LATITUDE,LONGITUDE,D_19000101)]

### GET PET ###
setkey(DT,PIXELID)
DT.temp <- DT[!is.na(D_19000101),]


dt.spei[,c("PIXELID","LATITUDE","LONGITUDE"):=DT.temp[,list(PIXELID,LATITUDE,LONGITUDE)]]

setkey(dt.spei,'PIXELID')

# Join to make maps
dt.spei.all <- DT[dt.spei]
dt.spei.all[,LATITUDE:=i.LATITUDE]
dt.spei.all[,LONGITUDE:=i.LONGITUDE]
dt.spei.all[,c('V1','V2','V3','V4','V5','V6','V7','V8','V9','V10','V11','V12'):=NULL]



MapPixels <- function (DT_Pixels, vraster.template, vext) {
  # Generates a raster based on x,y pixels.
  #
  # Args:
  #   DT_Pixels: A data tables with x,y coordinates and the values to assign each cell
  #   vraster.template: A string value with the filename of a raster, this
  #                     raster is used to get the CRS information
  #   vext: This should contain the "extent" of the analysis area
  # Returns:
  #   A raster of the input pixels

  if (nchar(vraster.template) <= 1) {
    #vraster.template <- tk_choose.files(caption="Select raster template...")
    vraster.template <- raster(vraster.template)
    vraster.template <- crop (vraster.template, vext)
  } else {
    vraster.template <- raster(vraster.template)
    #vraster.template <- crop(vraster.template, vext)
  }
  coordinates(DT_Pixels) <- ~ x + y
  gridded(DT_Pixels)<- TRUE
  vRasterOut <- raster(DT_Pixels)
  crs(vRasterOut) <- crs(vraster.template)

  return (vRasterOut)
}

vRaster <- '~/smbdir/PRISM800m/MONTHLY/ppt/cai_ppt_us_us_30s_193712.tif'
vext <- extent(raster(vRaster))

ptm<-proc.time()
st <-seq(as.Date("1901/1/1"),as.Date("2013-12-01"), by = "month")
for (j in seq_len(ncol(dt.spei.all)-3)) {
	#seq_len(ncol(dt))) {
	map <- MapPixels(dt.spei.all[,list(x=LONGITUDE,y=LATITUDE,spei=dt.spei.all[[j]])], vRaster, vext)
	vname <- paste0("SPEI_CONUS_",as.character(st[j]),".tif")
	writeRaster(map,vname,format='GTiff', datatype = 'FLT4S', NAflag=-9999)
}
proc.time() - ptm
#