# fields, Tools for spatial data
# Copyright 2004-2011, Institute for Mathematics Applied Geosciences
# University Corporation for Atmospheric Research
# Licensed under the GPL -- www.gpl.org/licenses/gpl.html
`Wendland.father` <-
function (x, theta = 1, dimension = 1, k=3) 
{
  Wendland(abs(x), theta=theta, dimension=dimension, k=k)}

`Wendland.mother` <-
function (x, theta = 1, dimension = 1, k=3) 
{ 
  sign(x)*Wendland(abs(x), theta=theta, dimension=dimension, k=k, derivative=1)}

"wendland.basis" <- function(x1, x2, theta = 1, V=NULL,
    k = 3, C = NA, Dist.args = list(method = "euclidean"), 
    spam.format = TRUE, verbose = FALSE, flavor=0) {
    
#  the rest of the possiblities require some computing
# setup the two matrices of locations
#
  d <- ncol(x1)
  if( d!=2) stop("only 2-d basis is supported")
  n1 <- nrow(x1)
  n2 <- nrow(x2)
# catch bad theta format
  if (length(theta) > 1) {
       stop("theta as a matrix or vector has been depreciated") }
  if (!is.null(V)) {
        if (theta != 1) {
            stop("can't specify both theta and V!")
        }
        x1 <- x1 %*% t(solve(V))
        x2 <- x2 %*% t(solve(V))
    }
  delta <- 1.5*theta 
# once scaling is done taper is applied with default range of 1.0
#  d dimension and  k is the order
#  first  find sparse matrix of Euclidean distances
#  ||x1-x2||**2 (or any other distance that may be specified by
#   the method component in Dist.args
#  any distance beyond delta is set to zero -- anticipating the
# tapering to zero by the Wendland.
#
  sM <- do.call("nearest.dist", c(list(x1, x2, delta = delta, 
        upper = NULL), Dist.args))
# scale distances by theta
# note: if V is passed then theta==1 and all the scaling should be done with the V matrix.
# there are two possible actions listed below:
# find  Wendland cross covariance matrix
# return either in sparse or matrix format
# Create rowindices vector
  sMrowindices <- rep(1:n1, diff(sM@rowpointers))
  dx <- (x1[sMrowindices, 1] - x2[sM@colindices, 1])
  dy <- (x1[sMrowindices, 2] - x2[sM@colindices, 2])
# switches for different flavors    
  if(flavor==0){
        sM@entries <- Wendland.father(dx,theta,k=k)*
                          Wendland.father(dy,theta,k=k)}
  if(flavor==1){
        sM@entries <- Wendland.father(dx,theta,k=k)*
                          Wendland.mother(dy,theta,k=k)}
  if(flavor==2){
        sM@entries <- Wendland.mother(dx,theta,k=k)*
                          Wendland.father(dy,theta,k=k)}
  if(flavor==3){
        sM@entries <- Wendland.mother(dx,theta,k=k)*
                          Wendland.mother(dy,theta,k=k)}
# convert to full matrix format     
  if (!spam.format) {
          sM<- as.matrix(sM)}
# 
  if (is.na(C[1])) {
      return(sM)}
    else{
      return( sM%*%C)} 
       
    # should not get here!
}



LatticeKrig<- function( x,y=NULL, weights = rep(1, nrow(x)), NC=16, center.range=NULL, beta=-.2, lambda=1.0,
                       iseed=123,NtrA=20,
                          overlap=2.5, setup=FALSE){
  N<- nrow(x)
# if center range is missing use the locations
  if( is.null(center.range)){
  x1<- range( x[,1])
  x2<- range( x[,2])}
  else{
  x1<- center.range$x
  x2<- center.range$y}
# spacing for grid
  d1<- x1[2] - x1[1]
  d2<- x2[2] - x2[1]
# actual number of grid points is set so that centers are equally
# spaced in both axes. NC is the maximum number of grid points
# delta used to determine support of wendland basis functions
# default support is (-1,1)  -- i.e. if delta is 1.0 
  delta<- max(c(d1,d2))/NC
  center.grid.list<-  list( x= seq( x1[1], x1[2],delta ), y= seq( x2[1], x2[2],delta))
  N1<- length( center.grid.list$x)
  N2<- length( center.grid.list$y)
  np<- N1*N2
# create centers there should be N1*N2 of these
  centers<- make.surface.grid( center.grid.list)
  basis.delta<- overlap*delta
# Matrix of N1*N2 basis function (columns) evaluated at the N locations (rows)
# this can be a large matrix if not encoded in sparse format.   
  PHI<-  wendland.basis( x, centers, theta=basis.delta, spam.format=TRUE, flavor=0)
# square root of precision matrix of the lattice process
#   solve(t(H)%*%H) is proportional to the covariance matrix of the Markov Random Field
  H<-  lattice.precision(N1,N2, spam.format=TRUE, beta=beta)
  
# variational matrix used to find coefficients of fitted surface and evaluate
  
############################################################################################
# this is the regularized regression matrix that is the key to the entire algorithm:
########################################################################################  
  temp<- t(PHI)%*%PHI + lambda*(t(H)%*%H)
  nzero <- length(temp@entries)
#  
# S-M-W identity can be used to evaluate the data covariance matrix:
#   M =  PHI%*% solve(t(H)%*%H) %*% t( PHI) + diag( lambda, N)
# i.e. because temp is sparse, sparse computations can be used to indirectly solve
# linear systems based on M
############################################################################################  
# find Cholesky square root of this matrix
############################################################################################  
#  This is where the heavy lifting happens! Note that temp is sparse format so
#  implied by overloading is a sparse cholesky decomposition. 
#  if this has been coded efficiently this step should dominate all other computations. 
  chol( temp)-> Mc
  if( setup){
     return(
            list(centers=centers,  center.grid.list= center.grid.list, N1=N1, N2=N2,
              overlap=overlap, delta=delta,basis.delta=basis.delta,            
              Mc=Mc, np=np, lambda.fixed=lambda,nonzero.entries=nzero,
                 PHI=PHI, XtX=t(PHI)%*%PHI)
            )
   }
#
# find log determinant of reg matrix temp for use in the log likeihood
  lnDetReg <- 2 * sum(log(diag(Mc)))
  lnDetQ<-  2* sum( log( diag( chol(t(H)%*%H))))
# now apply a miraculous determinant identity (Sylvester''s theorem)
#  det( I + UV) = det( I + VU)    providing UV is square
# or more generally   
#  det( A + UV) = det(A) det( I + V A^{-1}U)
#  or as we use it
#  ln(det( I + V A^{-1}U)) = ln( det(  A + UV)) - ln( det(A))
#  
#  with A = lambda* t(H)%*%H , U= t(PHI) V= PHI   
#  M =  lambda*(I + V A^{-1}U)
#
#   lnDet (M) = N* ln(lambda) + ln(Det(I + V A^{-1}U))
#             = N* ln(lambda) + ln(Det( temp)) - ln Det( lambda*Q)
#             = N* ln(lambda) + ln(Det( temp))  - ln Det( Q) - N1*N2* log( lambda)
# Here we  canceled out lambda terms 
# 
  lnDetCov<- lnDetReg - lnDetQ + (N - N1*N2)* log(lambda)
# Amazing!
# 
# Spatial drift matrix -- assumed to be linear in coordinates. (m=2)
  m=2
  nt= 3
  T.matrix<- cbind( rep(1, N), x)
#
# The goal is   d.coef = (T^t M^{-1} T)^{-1} T^t %*% M^{-1} y
#  M  is  as defined in comments above
#  apply the inverse of M to T and y  using S-M-W formula and cholesky of
#  regularization matrix (Mc in code above)
#  
  A <- forwardsolve(Mc, transpose = TRUE, t(PHI)%*%T.matrix, upper.tri = TRUE)
  A <- backsolve(Mc, A)  
  A<- t(T.matrix)%*%(T.matrix - PHI%*%A)/lambda
#   A is  (T^t M^{-1} T) 
  b <- forwardsolve(Mc, transpose = TRUE, t(PHI)%*%y, upper.tri = TRUE)
  b <- backsolve(Mc, b)  
  b<- t(T.matrix)%*%(y- PHI%*%b)/lambda
# b is   (T^t M^{-1} y)
# Save the intermediate matrix   (T^t M^{-1} T) ^{-1}
# this the GLS covariance matrix of estimated coefficients
# should be small -- the default is 3X3  
  Omega<- solve( A)
# GLS estimates   
  d.coef<- Omega%*% b

# coefficients of basis functions.   
  c.coef <- forwardsolve(Mc, transpose = TRUE, t(PHI)%*%(y-T.matrix%*%d.coef),
                         upper.tri = TRUE)
  c.coef <- backsolve(Mc, c.coef)
  fitted.values<- T.matrix%*%d.coef + PHI%*%c.coef
  residuals<- y- fitted.values
  c.mKrig<- (y- fitted.values)/lambda

# finding quadratic form
# this uses a shortcut formula for the quadratic form in terms of
# the residuals 
    quad.form<-   sum(y* c.mKrig )
# MLE estimate of rho and sigma
# these are derived by assuming Y is  MN(  Td, rho*M )  
  rho.MLE <- quad.form/N
  shat.MLE <- sigma.MLE <- sqrt(lambda * rho.MLE)
# the  log profile likehood with  rhohat  and  dhat substituted
# leaving a profile for just lambda.
# note that this is _not_  -2*loglike just the log and
# includes the constants
  lnProfileLike <- (-N/2 - log(2*pi)*(N/2)
                      - (N/2)*log(rho.MLE) - (1/2) * lnDetCov)
# Monte Carlo estimate of effective degrees of freedom
# uses the fact that the expected value of a quadratic form is the
# trace of the matrix.
#  NtrA samples are taken and  typically NtrA= 20 is adequate
#  to use the sample mean as an estimate of the expectation.
#  A(lambda,beta)  here is of the course the smoothing matrix and predicted
#  values found the same way as the estimate. Note use of vectorization
#  to avoid a loop over the Monte Carlo samples. 
#
  set.seed(iseed)
 
  Ey<- matrix( rnorm( NtrA*N), N,NtrA)
  b <- forwardsolve(Mc, transpose = TRUE, t(PHI)%*%Ey, upper.tri = TRUE)
  b <- backsolve(Mc, b)
  b<- t(T.matrix)%*%(Ey- PHI%*%b)/lambda
  d.temp<-  solve( A, b)
  
  c.temp <- forwardsolve(Mc, transpose = TRUE, t(PHI)%*%(Ey- T.matrix%*%d.temp),
                                         upper.tri = TRUE)
  c.temp <- backsolve(Mc, c.temp)
  
  Eyhat<-T.matrix%*%d.temp + PHI%*%c.temp 
  trA.info <- t(Ey *(Eyhat))%*%rep(1,N)
  trA.est <- mean(trA.info)

# find the GCV function
  GCV=  (sum((residuals)^2)/N )  /( 1- trA.est/np)^2
  
  obj<-list(x=x,y=y,d.coef=d.coef, c.coef=c.coef,
              fitted.values=fitted.values, residuals= residuals,
              centers=centers, beta=beta, center.grid.list= center.grid.list, N1=N1, N2=N2,
              overlap=overlap, delta=delta,basis.delta=basis.delta,
              GCV=GCV, lnProfileLike= lnProfileLike,
              rho.MLE=rho.MLE, shat.MLE= shat.MLE,lambda=lambda,
              lnDetCov= lnDetCov, quad.form= quad.form,
              Mc=Mc,
              trA.info=trA.info, trA.est=trA.est,
              eff.df= trA.est, np=np, lambda.fixed=lambda,nonzero.entries=nzero,m=m,nt=nt,
              call=match.call()
            )
  class(obj)<- "LatticeKrig"
  return(obj)
}

surface.LatticeKrig<- function(obj,...){
  surface.Krig( obj,...)}

predict.LatticeKrig<- function( object, xnew=NULL,...){
  if( is.null(xnew)){
    xnew<- object$x}
  PHIg<-  wendland.basis( xnew, object$centers, theta=object$basis.delta, spam.format=TRUE, flavor=0)
  NG<- nrow( xnew)
  T.matrix.g<- cbind( rep(1,NG), xnew)
return( T.matrix.g%*% object$d.coef + PHIg%*%object$c.coef)
}

print.LatticeKrig <- function(x, digits=4, ...){
   if (is.matrix(x$residuals)) {
        n <- nrow(x$residuals)
        NData <- ncol(x$residuals)
    }
    else {
        n <- length(x$residuals)
        NData <- 1
    }
    c1 <- "Number of Observations:"
    c2 <- n
    if (NData > 1) {
        c1 <- c(c1, "Number of data sets fit:")
        c2 <- c(c2, NData)
    }
    c1 <- c(c1, "Degree of polynomial null space ( base model):")
    c2 <- c(c2, x$m - 1)
    c1 <- c(c1, "Number of parameters in the null space")
    c2 <- c(c2, x$nt)
    if (!is.na(x$eff.df)) {
        c1 <- c(c1, " Eff. degrees of freedom")
        c2 <- c(c2, signif(x$eff.df, digits))
        if (length(x$trA.info) < x$np) {
            c1 <- c(c1, "   Standard Error of estimate: ")
            c2 <- c(c2, signif(sd(x$trA.info)/sqrt(length(x$trA.info)), 
                digits))
        }
    }
    c1 <- c(c1, "Smoothing parameter")
    c2 <- c(c2, signif(x$lambda.fixed, digits))
    if (NData == 1) {
        c1 <- c(c1, "MLE sigma ")
        c2 <- c(c2, signif(x$shat.MLE, digits))
        c1 <- c(c1, "MLE rho")
        c2 <- c(c2, signif(x$rho.MLE, digits))
    }
    c1 <- c(c1, "Nonzero entries in covariance")
    c2 <- c(c2, x$nonzero.entries)
    sum <- cbind(c1, c2)
    dimnames(sum) <- list(rep("", dim(sum)[1]), rep("", dim(sum)[2]))
    cat("Call:\n")
    dput(x$call)
    print(sum, quote = FALSE)
    cat(" ", fill = TRUE)
    cat("Covariance Model: Wendland/Lattice", fill = TRUE)
    cat("Lattice is ", x$N1,"X",x$N2, fill=TRUE)
    cat( "total number of points: ", x$N1*x$N2, " with spacing", x$delta,  fill=TRUE)
    cat("Value for lattice dependence (beta): ", x$beta,fill = TRUE)
    invisible(x)
}                      

"lattice.precision" <- function(N1,N2, spam.format = TRUE, beta) {
  NN<- N1*N2
  da<- as.integer( c(NN,NN))
  I<- as.integer(rep(1:N1,N2))
  J<- as.integer(rep((1:N2), rep( N1,N2)))
# contents of sparse matrix
  ra<- rep( 1, N1*N2)
  ra<- c( ra, rep( beta, 4*NN))
#  Note: if beta depends on lattice position
# pass beta as a matrix and use  ra<- as.numeric(c( ra,rep( c(beta), 4)))
  Hi<- rep(1:NN,5) 
  Hj<- c( I + (J-1)*N1,
         (I-1) + (J-1)*N1, (I+1) + (J-1)*N1,
         I + (J-2)*N1, I + (J)*N1 )   
  good<- c( rep( TRUE,NN),
          (I-1) > 0, (I<N1),
          (J-2)>=0, (J<N2))
# remove cases that are beyond the lattice (edges, corners)        
  Hi<- as.integer(Hi[good])
  Hj<- as.integer(Hj[good])
  ra<- ra[good]
# Convert matrix to spam format  
  sM <- spind2spam( list( ind=cbind(Hi,Hj), ra=ra, da=da) )
# convert to an ordinary matrix (usually for debugging)     
  if (!spam.format) {
          sM<- as.matrix(sM)}
  return(sM)
}

LatticeKrig.cov<- function( x1,x2=NULL,center.grid.list, beta=-.2, overlap=2.5, C=NA, marginal=FALSE){
  if( is.null(x2)){ x2<-x1}
  N1<- length( center.grid.list$x)
  N2<- length( center.grid.list$y)
  centers<- make.surface.grid( center.grid.list)
  delta<- max(c( center.grid.list$x[2] - center.grid.list$x[1],
              center.grid.list$y[2] - center.grid.list$y[1] ))
  basis.delta<- overlap*delta
  PHI1<-  wendland.basis( x1, centers, theta=basis.delta, spam.format=FALSE, flavor=0)
  PHI2<-  wendland.basis( x2, centers, theta=basis.delta, spam.format=FALSE, flavor=0)
  H<-  lattice.precision(N1,N2, spam.format=TRUE, beta=beta)
#  Note:  Q  <- t(H)%*%(H)
  chol(t(H)%*%(H))-> Mc 
  A <- forwardsolve(Mc, transpose = TRUE, t(PHI2), upper.tri = TRUE)
  A <- backsolve(Mc, A)
  
   if (is.na(C[1]) & !marginal) {
        return( PHI1%*%A)
    }
    if (!is.na(C[1])) {
        return(PHI1%*%A %*% C)
    }
    if (marginal) {
        return(rep(1, nrow(x1)))
    }

 
}