library(maps)
library(akima)
library(fields)

nrows=72
ncols=36
ntime = 118    #Nov-Mar 1901 - Nov-Mar 2018

nyrs = ntime    #1902 - 2016

nglobe = nrows*ncols
N = nrows*ncols


### Lat - Long grid..

locs=matrix(scan("http://civil.colorado.edu/~balajir/CVEN6833/R-sessions/session2/files-4HW2/sst-lat-long.txt"), ncol=2, byrow=T)
ygrid=seq(-87.5,87.5,by=5)
ny=length(ygrid)

xgrid=seq(27.5,382.5,by=5)
#xgrid[xgrid > 180]=xgrid[xgrid > 180]-360	#longitude on 0-360 grid if needed
xgrid[xgrid > 180]=xgrid[xgrid > 180]
nx=length(xgrid)

xygrid=matrix(0,nrow=nx*ny,ncol=2)

i=0
for(iy in 1:ny){
for(ix in 1:nx){
i=i+1
xygrid[i,1]=ygrid[iy]
xygrid[i,2]=xgrid[ix]
}

}

# REad Kaplan SST data..


data=readBin("http://civil.colorado.edu/%7Ebalajir/CVEN6833/HWs/HW-2-2018/Kaplan-SST-NDJFM1900-NDJFM2018.r4",what="numeric", n=( nrows * ncols * ntime), size=4,endian="swap")


data <- array(data = data, dim=c( nrows, ncols, ntime ) )


data1=data[,,1]
	

# the lat -long data grid..

index=1:(nx*ny)

index1=index[data1 < 20 & data1 != "NaN"]	# only non-missing data.
xygrid1=xygrid[index1,]

x1=xygrid1[,2]
#x1[x1 < 0]= x1[x1 < 0] + 360
#xygrid1[,2]=x1

nsites=length(index1)
data2=data1[index1]

### SSTdata matrix - rows are seasonal (i.e. one value per year)
## and columns are locations
sstdata=matrix(NA,nrow=nyrs, ncol=nsites)


for(i in 1:nyrs){
data1=data[,,i]
index1=index[data1 < 20 & data1 != "NaN"]
data2=data1[index1]
sstdata[i,]=data2
}

sstannavg = sstdata
indexgrid = index1
rm("data")	#remove the object data to clear up space


## write out the grid locations..
write(t(xygrid1),file="kaplan-sst-locs.txt",ncol=2)


###################### PCA 
##

#get variance matrix..
## scale the data
sstscale = scale(sstannavg)
zs=var(sstscale)

#do an Eigen decomposition..
zsvd=svd(zs)

#Principal Components...
pcs=t(t(zsvd$u) %*% t(sstscale))

#Eigen Values.. - fraction variance 
lambdas=(zsvd$d/sum(zsvd$d))

plot(1:25, lambdas[1:25], type="b", xlab="Modes", ylab="Frac. Var. explained")
grid()


#plots..
#plot the first spatial component or Eigen Vector pattern..


# the data is on a grid so fill the entire globaal grid with NaN and then populate the ocean grids with
# the Eigen vector
xlong = sort(unique(xygrid[,2]))
ylat = sort(unique(xygrid[,1]))
zfull = rep(NaN,nglobe)   #also equal 72*36
zfull[indexgrid]=zsvd$u[,1]
zmat = matrix(zfull,nrow=nrows,ncol=ncols)

image.plot(xlong,ylat,zmat,ylim=range(-40,70))
contour(xlong,ylat,(zmat),ylim=range(-40,70),add=TRUE,nlev=6,lwd=2)
map("world2",add=T)
#world(add=TRUE,shift=TRUE)


### Similarly plot the other three Eigen vectors..


plot(1901:2018, pcs[,1],type="b",xlab="Year",ylab="PC1")

## similarly plot other PCs - i.e. temporal modes

## ENSO index
nino4=scan("http://iridl.ldeo.columbia.edu/SOURCES/.Indices/.nino/.EXTENDED/.NINO4/T/(Nov%201899)/(Mar%202018)/RANGE/T/(Nov-Mar)/seasonalAverage/data.ch")


## plot the PCs
## PC1 will be a trend
plot(1901:2018, scale(pcs[,1]),type="b",xlab="Year",ylab="PC1")

## PC2 will be ENSO
plot(1901:2018, scale(pcs[,2]),type="b",xlab="Year",ylab="PC1")
lines(1901:2018, nino4, col="red")


## similarly plot PCs 2,3 

## Note that cor(nino34, pcs[,2]) gives the highest correlation
##meaning the second mode is ENSO

## also the first PC has a strong trend implies that it is the
#global warming pattern!


################## If you wish to remove a component from the data, say first component.
####

nmodes = length(zsvd$u[1,])   # number of modes
nkeep = c(1)		# modes to keep, here we keep the first mode. If more then
			# nkeep=c(1,2,3) etc..
E = matrix(0,nrow=nmodes,ncol=nmodes)
E[,nkeep]=zsvd$u[,nkeep]
sstannavgkeep = pcs %*% t(E)

sstanrem = sstscale - sstanavgkeep

## Now perform PCA on sstanrem


########### (ii)
###  Rotate first six PCS
##################
zrot = varimax(zsvd$u[,1:6],normalize=FALSE)
#zrot = promax(zsvd$u[,1:6],m=2)

## plot the first rotated spatial mode
xlong = sort(unique(xygrid[,2]))
ylat = sort(unique(xygrid[,1]))
zfull = rep(NaN,nglobe)   #also equal 72*36
zfull[indexgrid]=zrot$loadings[,3]
zmat = matrix(zfull,nrow=nrows,ncol=ncols)

image.plot(xlong,ylat,zmat,ylim=range(-40,70))
contour(xlong,ylat,(zmat),ylim=range(-40,70),add=TRUE,nlev=6,lwd=2)
map("world2",add=T)
#world(add=TRUE,shift=TRUE)


### Similarly plot the other three rotated Eigen vectors..

rotpcs=t(t(zrot$loadings) %*% t(sstannavg))
## plot the rotated PCs
## PC1 will be a trend
plot(1901:2018, scale(rotpcs[,1]),type="b",xlab="Year",ylab="PC1")

## PC2 will be ENSO
plot(1901:2018, scale(rotpcs[,2]),type="b",xlab="Year",ylab="PC1")
lines(1901:2018, nino4, col="red")


## similarly plot rotated PCs 2,3