# Apply the EnviroStat package version of the early WLS deformation algorithm
library(EnviroStat)
library(MASS)

source("deformloglik.l1.R")
source("visualize.tps.warps.R")
source("draw2.R")  # to replace a function in EnviroStat

# For reading the French precipitation dataset
plot.title <- "Pluie.LanguedocRoussillon"
setwd("/Users/pdsampson/Dropbox/PASI/PDS Nonstationary covariances/Practicum")
Pluie <- read.csv("pluieNA.novdec.csv")
Pluie <- Pluie[,-1]
PluieStations <- read.csv("pluieStations.Practicum.csv")
X <- as.matrix(PluieStations[,c("x","y")])
S <- cov(log(Pluie+1),use="pair")  # Pairwise covariances for log precip
t <- nrow(Pluie)

X0 <- X  # Save initial configuration      
n <- nrow(X)  # number of spatial points  

# Standardize the coordinate configuration by centering to mean zero,
# scaling to "centroid size" one, and rotating to the principal 
# axes of the configuration.
Xs <- scale(X,scale=FALSE)  
Xs <- Xs/sqrt(sum(Xs^2))
Xsvd <- svd(Xs)
# Here I will make sure diag elements of rotation matrix are positive
# by scaling columns by +/- 1.  This is to keep the rotation from flipping
# the orientation entirely.
Xsvd$v[,1] <- Xsvd$v[,1] * sign(Xsvd$v[1,1]) # 
Xsvd$v[,2] <- Xsvd$v[,2] * sign(Xsvd$v[2,2])
Xt <- Xs %*% Xsvd$v
xscaling <- list(mean=attr(Xs,"scaled:center"),
                 #scale=attr(Xs,"scaled:scale")*sqrt(2*(n-1)),
                 scale=1,
                 rotation=Xsvd$v)
X0 <- X   # Save initial coordinates
X <- Xt   # Replace X with center, scaled, rotated coordinates


### 1.  Fitting using functions from "EnviroStat" package

# "Dispersion" based on correlation matrix
Cor <- S/sqrt(diag(S)%o%diag(S))
Disp <- 2 - 2*Cor
Gdist <- as.matrix(dist(X,upper=T,diag=T))
plot(Gdist,Disp)

# Fit an exponential variogram for the dispersion data
h.lt <- Gdist[row(Gdist) < col(Gdist)]
disp.lt <- Disp[row(Disp) < col(Disp)]
variogfit <- Fvariogfit3(disp.lt, h.lt, a0 = 0.1, t0 = 1, verbose=T)
x <- seq(min(h.lt), max(h.lt), .01)
a0 <- variogfit$a[1]
t0 <- variogfit$t0

plot(-.2, 0, xlim = c(0, max(h.lt)), ylim = c(0, 2),
  xlab = "Dist", ylab = "Dispersion", type = "n",
  main = c("Intersite dispersion vs intersite distance"))
for (i in 1:(n-1))
  for(j in (i+1):n)
    points(Gdist[i, j], Disp[i, j])
lines(x, a0 + (2 - a0) * (1 - exp(-(t0 * x))),col="red")

# Step 1.  Alternating estimation of variogram and deformation 
# (first with no constraints on the deformation)

sg.est <- Falternate3(Disp, X, max.iter = 100,
   alter.lim = 25, model = 1, lambda=0)   
                 
# Step 2.  Select a smoothing parameter interactively
#apply(coords.lamb, 2, range)
coords.grid <- Fmgrid(range(X[,1]),
                        + range(X[,2]))
par(mfrow = c(1, 2))
temp <- setplot(X, axes = TRUE)
deform <- Ftransdraw(disp = Disp, Gcrds = X,
  MDScrds = sg.est$ncoords,
  gridstr = coords.grid)

# Step 2b.  Re-estimate variogram and deformation with
# specified smoothing.
sg.est <- Falternate3(Disp, X, max.iter = 100,
  alter.lim = 25, model = 1, lambda=0.01)   

# Step 3.
# Compute the thin-plate spline from the final result of Ftransdraw
Tspline <- sinterp(X, deform$Dcrds, lam = 0 )
temp <- Ftransdraw(disp=Disp,Gcrds=X,
                   MDScrds=Tspline$y, 
                   gridstr=coords.grid)

# Draw a "biorthogonal grid" illustrating the spatially varying
# affine derivative (local anisotropy) of the deformation or 
# correlation structure.
par(mfrow = c(1, 1))
Tgrid <- bgrid(start = c(0, 0), xmat = X,
    coef = Tspline$sol)
tempplot <- setplot(X, axes = TRUE)
text(X, labels = 1:nrow(X))
draw2(Tgrid, fs = TRUE)



###  2.  Now switch to fitting with partial warps
###  Lots of different notation from above

# Reasonable initial values for parameters of an
# exponential correlation model
theta     <- log(.2)   # range parameter
sigmaSqrd <-  log(.05)   # scale parameter
eta       <-  log(.1)   # nugget parameter
spatialParams <- c(theta, sigmaSqrd, eta)            

# Initialize matrix of partial warp coefficients:
# -  the first column consists of the 'a' and 'b" affine parameters.
# -- the last n-3 columns are coefficients of the partial warps
# for the partial warps
beta0 <- matrix(0,nrow=2,ncol=n-2)  

## calculating bending energy matrix B;
# See the texts by Bookstein (1991) or Dryden and Mardia (1998)

h <- as.matrix(dist(X,diag=TRUE,upper=TRUE))
K <- h^2 * log(h)
diag(K) <- 0
one <- rep(1,n)
Gamma <- cbind(K,one,X)
Gamma <- rbind(Gamma,c(one,0,0,0))
Gamma <- rbind(Gamma,cbind(t(X),matrix(0,2,3)))
Ginv <- solve(Gamma)
Ginv <- (Ginv + t(Ginv))/2  # make exactly symmetric prior to eigen

B <- Ginv[1:n,1:n]
Beig <- eigen(B)
g <- Beig$vectors
l <- Beig$values

# *** Need two versions of deformloglik: (1) one for estimation of correlation 
# parameters given the deformation parameters, "beta", and (2) one for
# estimation of deformation parameters given the correlation parameters.
# 
# In addition to specifying which are the parameters ("par") to be optimized,
# the 'deformloglik.beta' function has arguments to allow specification
# of subset of the partial warps to be included in the fitting.
# 'jWarp' is a vector specifying which of the partial warps to include.
# Generally initialize all the partial warp parameters to zero, meaning
# we start with an identity mapping.

jWarp <- 1:(n-3)
# Or selext a subset of warps, throwing out those with
# high bending energy.  E.g.
jWarp <- 1:10
lambda <- 1.5  # L1 constraint

beta0 <- matrix(0,nrow=2,ncol=1+length(jWarp))
beta <- as.vector(beta0)

### 3.  Alternating optimization of variogram parameters and coefficients
###     for partial warps.  
###     This version runs a specified number of iterations; is not set up 
###     to assess convergence.

#nrep <- 25
nrep <- 5

for (i in 1:nrep) {
  # (1) Optimize variogram parameters
  beta <- matrix(beta,nrow=2,ncol=1+length(jWarp))
  betafull <- matrix(0,nrow=2,ncol=(n-2))
  betafull[,c(1,jWarp+1)] <- beta
  likfit.variog <- optim(par=spatialParams,deformloglik.variog.lambda,hessian=T,
                         B=B,X=X,beta=betafull,S=S,t=t,lambda=lambda)  #,
  #method="BFGS")
  spatialParams <- likfit.variog$par
  print(round(exp(spatialParams),4))
  print(paste("convergence",likfit.variog$convergence))
  # (2) Optimize deformation parameters
  beta <- as.vector(beta)
  #likfit.beta <- optim(par=beta,deformloglik.beta,hessian=T,
  #                   control=list(maxit=500),B=B,X=X,
  #                   spatialParams=spatialParams,S=S,t=t,jWarp=jWarp,
  #                   method="BFGS")
  likfit.beta <- optim(par=beta,deformloglik.beta.lambda,hessian=T,
                       control=list(maxit=500,trace=1),B=B,X=X,
                       spatialParams=spatialParams,S=S,t=t,jWarp=jWarp,lambda=lambda,
                       method="BFGS")
  beta <- likfit.beta$par
  
  betafit <- matrix(likfit.beta$par,nrow=2)
  print(round(betafit,2))
  print(paste("convergence",likfit.beta$convergence))
}

betafitfull <- matrix(0,nrow=2,ncol=(n-2))
betafitfull[,c(1,jWarp+1)] <- betafit

theta <- exp(spatialParams[1])
sigmaSqrd <- exp(spatialParams[2])
eta <- exp(spatialParams[3])

### 4. Plot results, including
###   (a) the deformation and plots of covariance and correlation vs distance
###   (b) visualization of contribution of each of the partial warp components
###   (c) visualization of each of the partial warps, separately for x and y 
###       components, arbitrarily scaled.

## (a)

par(mfrow=c(1,2))
Fdraw_deform(X,beta=betafitfull,xscaling=xscaling,colpts="black",cexpts=.5)
Y <- Fcomp_deform(X,beta=betafitfull,X)

##
fitted.warps <- "all"  # or, for example,
fitted.warps <- "1:10"
plot.title <- "Pluie"
pdftitle1 <- paste(plot.title,", lambda=",lambda,", warps ", fitted.warps," .pdf",sep="")
pdf(pdftitle1)
#pdf("test")

par(mgp=c(2,.75,0),mar=c(3,2,2,1))
eqscplot(X0[,1],X0[,2],type="n",tol=.25)
text(X0[,1],X0[,2],cex=.75)
title(main="Lambert projection coords")
par(mfrow=c(2,1))       
Fdraw_deform(X,beta=betafitfull, xscaling=xscaling,cexpts=.75)

par(mfrow=c(2,2),mar=c(4,4,3,1))
Gdist <- as.matrix(dist(X,upper=T,diag=T))
plot(Gdist,S)
Dsort <- sort(Gdist)
Esort <- sigmaSqrd * exp( (-1/theta) * Dsort)
Esort[Dsort==0] <- Esort[Dsort==0] + eta
lines(Dsort,Esort,col="red")

Ddist <- as.matrix(dist(Y,upper=T,diag=T))
plot(Ddist,S)
Dsort <- sort(Ddist)
Esort <- sigmaSqrd * exp( (-1/theta) * Dsort)
Esort[Dsort==0] <- Esort[Dsort==0] + eta
lines(Dsort,Esort,col="red")

Cor <- S/sqrt(diag(S)%o%diag(S))
plot(Gdist,Cor)       
Ecorsort <- sigmaSqrd * exp( (-1/theta) * Dsort)/(eta+sigmaSqrd)
lines(Dsort,Ecorsort,col="red")
plot(Ddist,Cor)
lines(Dsort,Ecorsort,col="red")

dev.off()

## (b)
# Visualize each of the component pieces of the fit.
# (b1) Visualize each of the warps to an arbitrary scale
# (b2) Visualize the actual partial warps.  *** TO BE DONE ***
pdftitle2 <- paste(plot.title,", lambda=",lambda,", by warp ", fitted.warps," .pdf",sep="")
pdf(pdftitle2)

par(mfrow=c(2,2))
# Affine
beta <- betafitfull
beta[,2:(n-2)] <- 0
Fdraw_deform(Xt,beta,cexpts=.75,drawsq=F)
title(main=paste("Affine, coefs = ",paste(round(beta[,1],2),collapse=",")))

for (j in (1+jWarp)) {
  beta <- betafitfull
  beta[,(1:(n-2))!=j] <- 0
  Fdraw_deform(Xt,beta,cexpts=.75,drawsq=F)
  title(main=paste("PW",j,paste(round(beta[,j],2),collapse=",")))
}
dev.off()

# (c)  Draw biorthogonal grid
pdf(paste(plot.title,", lambda=",lambda,", warps ", fitted.warps," bgrid.pdf",sep=""))
par(mgp=c(2,.75,0))
Tspline <- sinterp(X, Y, lam = 0 )
Tgrid <- bgrid(start = c(0, 0), xmat = X,
               coef = Tspline$sol, perpc=10)
tempplot <- setplot(X, axes = TRUE)
text(X, labels = 1:nrow(X),cex=.75)
Bgrid.title <- paste("Bgrid: ",plot.title,"\n lambda=",lambda,", warps ", fitted.warps,sep="")
temp <- draw2(Tgrid, fs = TRUE, lwidth=c(2,2), lcolor=topo.colors(5), legend=TRUE)  #,lwidth=c(1,3))
title(main=Bgrid.title)
title(sub="Lines are colored by directional magnitude of the
      principal axes of the deformation")
dev.off()