# This is a script to calculate SPEI02 for 6 LOCA models x 2 RCPs
# Observations: 1/8 x 1/8 degree, 1950-2099
#===================================================================================================================
library(SPEI)
#===================================================================================================================

spei.ind = 2
tstart="1950-01"; tend="2099-12";
years=c(1950:2099)
scenarios=c('rcp45','rcp85')
models = c('canesm2','ccsm4','gfdl-cm3','giss-e2-r','hadgem2-es','miroc5')

for (n in 1:length(scenarios)){
  print(scenarios[n])
  dir = paste('./data/',scenarios[n],'/',sep='')
  files=(Sys.glob(paste(dir,'pr/*.Rdata', sep='')))
  files2=(Sys.glob(paste(dir,'pet/*.Rdata', sep='')))

  for (k in 1:length(models)){  
    fin = files[k]; load(fin)
    fmodel = unlist(strsplit(fin,"[._]"))[3]; print(fmodel)
    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[k])
    PET = aperm(ss$data, c(3,2,1)) # calculated by IEc using Modified-Hargreaves method
 
    #--- Calculate mm SPEI02
    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, 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(dir,'spei02/spei02_',fmodel,'_195001-209912.Rdata',sep=''))
    #image.plot(lon, lat, spei[,,3], zlim=c(-3,3), col=col.pal)
    rm(spei); print("completed")

  } # models k
} # scenarios n

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