###################################################
# A Standardized Precipitation Index for
# NCEP precipitation data http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanalysis.derived.surfaceflux.html
# Joseph Wheatley Biospherica March 2010
################################################

library("ncdf")

#data in prate.mon.mean.nc netcdf file
ncep.prec <- open.ncdf("prate.mon.mean.nc")
lons <- get.var.ncdf(ncep.prec,"lon")
lats <- get.var.ncdf(ncep.prec,"lat")
time <- get.var.ncdf(ncep.prec,"time")
prate <-get.var.ncdf(ncep.prec,"prate")#units Kg/m2/s
prate <- 2629800*prate #convert units mm/month
#dimensions
prate.dim <- dim(prate)
Nx <- prate.dim[1]
Ny <- prate.dim[2]
Nt <- prate.dim[3]
#flip upside-down Netcdf files
prate <- prate[,Ny:1,]
lats <- rev(lats)
#prate.mon is a new array containg year,month index
#preliminaries
dates<- seq(from=1948+1/24, to= 2011,by=1/12)
yeari <- function(i) {floor(i/12)+1} #
monthi <- function(i) {ifelse(i %% 12 == 0,12,i %% 12)}
years <- length(unique(sapply(seq(1:Nt),yeari)))
#initialise and fill prate.mon
prate.mon<- array(rep(NA,Nx*Ny*years*12),dim=c(Nx,Ny,years,12))
for(t in 1:Nt){ prate.mon[,,yeari(t),monthi(t)] <- prate[,,t]}

#data for selected location
#################################
longitude <- 145 #range 0, 360
latitude <- -30 # range -90,90
#################################
#find data closest grid cell
pdata1 <- prate[which.min((lons-longitude)^2),which.min((lats-latitude)^2),]
dates1 <- seq(from=1948+1/24, by = 1/12,length.out=length(pdata1))
# is there a trend?
prec.trend <- lm(pdata1~dates1) #linear fit
pdata1.detrend <- pdata1 - coef(prec.trend)[1]-coef(prec.trend)[2]*dates1 #detrend
pdata1.detrend <- pdata1.detrend - min(pdata1.detrend) # shift to remove any negative values

#pdata[[m]] contains average monthly precip over m months including current
pdata <- list(rep(0,12)) #initialise pdata
pdata[[1]] <- pdata1.detrend
for(m in 2:12){
pdata[[m]] <- sapply( seq(from=1, to=length(pdata1)),function(i) {im <- max(i-m+1,1); mean(pdata1.detrend[im:i])})
#add NAs at end so that all pdata[[m]] have same length
if(length(pdata[[m]]) < length(pdata)) pdata[[m]] <- c(pdata[[m]],NA)}
#pdata.mon[[m]] rewrites pdata[[m]] with a year, month index
pdata.mon <- list(rep(0,12))
for(i in 1:12){
pdata.mon[[i]] <- array(rep(NA,years*12),dim=c(years,12))
for(t in 1:Nt) pdata.mon[[i]][yeari(t),monthi(t)] <- pdata[[i]][t]}

#spi.m[[k]] is the spi index for timescale k
#time-scale loop. k is timescale index 1 .. 12
spi.m <- as.list(rep(0,N_timescale))
for (k in 1:12){
spi.m[[k]] <- as.list(rep(0,12))
# month loop. i is month index
for (i in 1:12)
{
  data <- na.omit(pdata.mon[[k]][,i])
  fit.cdf <- ecdf(data)
  cdfs <- sapply(data,fit.cdf) #locate on dgamma cdf - non-exceedence probablity
  spi.m[[k]][[i]] <- qnorm(cdfs) #inverse normal cdf
}
# add NA for missing months in current year
for (i in 1:12) if( length(spi.m[[k]][[i]])<years) spi.m[[k]][[i]] <- c(spi.m[[k]][[i]],NA)
#stitch the months together to form a time-series vector
spi.m[[k]] <- as.vector(t(matrix(unlist(spi.m[[k]]),ncol=12)))

}
#2D plot
spi.breaks <- c(-2.4,-2,-1.6,-1.3,-0.8,-0.5,0.5,0.8,1.3,1.6,2,2.4)
spi.cols <- colorRampPalette(c("darkred","red","yellow","white","green","blue","darkblue"),space="rgb")
spi <- matrix(unlist(spi.m),756,12) #convert list to matrix for plotting
spi[(spi==Inf)] <- 2.2 #necessary to remove Infs because ecdf is being used
#plot
filled.contour(dates,seq(1:12),spi,col=spi.cols(11),xlab="",ylab="time-scale (months)",cex.lab=1.7,font.axis=2,font.lab=2,levels=spi.breaks,key.title="SPI")
title(main="SPI",cex.main=2)