# This is a script to calculate SPEI02 from LOCA-obs data
# Observations: 1/8 x 1/8 degree, 1950-2013
#===================================================================================================================
library(SPEI)
#library(colorRamps)
#library(RColorBrewer)
#library(fields)
#library(maps)
#library(caTools)
#library(abind)
#===================================================================================================================

spei.ind = c(1,2,3)
tstart="1950-01"; tend="2013-12";
years=c(1950:2013)

dir = './data/LOCA.obs/'
  files=(Sys.glob(paste(dir,'*pr*.Rdata', sep='')))
  files2=(Sys.glob(paste(dir,'*pet*.Rdata', sep='')))
  
  load(files)
  PRCP = aperm(ss$data, c(3,2,1)) # convert to months x lat x lon
  lon=ss$lon; lat=ss$lat; time=ss$time

  load(files2)
  PET = aperm(ss$data, c(3,2,1)) # calculated by IEc using Modified-Hargreaves method
 
  #--- Calculate mm SPEI01-03
  for (k in 1:length(spei.ind)){
    print(paste(spei.ind[k],'-month SPEI',sep=''))
    wb = array(NA, dim=dim(PRCP))
    data = array(NA, dim=dim(PRCP))
  
    for (i in 1:length(lon)){
      print(i)
      wb[,,i] = PRCP[,,i] - PET[,,i] # calculate water balance
      data[,,i] = spei(ts(wb[,,i], freq=12, start=c(1950,1)), spei.ind[k], ref.start=c(1950, 1), ref.end=c(2005,12), na.rm=TRUE)$fitted
    }

  spei = aperm(data, c(3,2,1))
  spei[!is.finite(spei)] = NA # replace infinite values with NA
  res = list('data'=spei, 'lon'=lon, 'lat'=lat, 'time'=time)
  save(res, file=paste('./data/LOCA.obs/spei0',spei.ind[k],'_obs_195001-201312.Rdata',sep=''))
  #image.plot(lon, lat, spei[,,3], zlim=c(-3,3), col=col.pal)
  rm(spei); rm(res); i=1;
  print("completed")
}

#===================================================================================================================
#===================================================================================================================