# 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


LatticeKrig<- function( x,y=NULL, weights = rep(1, nrow(x)),Z=NULL, NC=16, center.range=NULL, beta=-.2, lambda=1.0,
                       iseed=123,NtrA=20,
                          overlap=2.5, setup=FALSE){
# make sure locations are a matrix and get the rows  
  x<- as.matrix(x)
  N<- nrow(x)
# make sure covariate is a matrix  
   if(!is.null(Z)){
      Z<- as.matrix(Z)}
# 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<- max(c(d1,d2))/(NC-1)
  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)
  
# delta used to determine support of wendland basis functions
  
  basis.delta<- overlap*delta
# weighted observation vector  
  wy<- sqrt(weights)*y
# Spatial drift matrix -- assumed to be linear in coordinates. (m=2)
# and includes Z covariate if not NULL  
  wT.matrix<- sqrt(weights)*cbind( rep(1, N), x,Z)
  nt<- ncol(wT.matrix)
  print(nt)
  nZ<- ifelse (is.null(Z), 0, ncol(Z))   
  ind.drift <- c(rep(TRUE, (nt-nZ)), rep( FALSE, nZ))
# Matrix of N1*N2 basis function (columns) evaluated at the N locations (rows)
# and multiplied by square root of diagonal weight matrix  
# this can be a large matrix if not encoded in sparse format.
  wS<-  diag.spam(sqrt( weights))
  wPHI<-  wS%*%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:
########################################################################################
  Q<- t(H)%*%H
  temp<- t(wPHI)%*%wPHI + lambda*(Q)
  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,
                 wPHI=wPHI, wT.matrix=wT.matrix,
                 weights=weights, nZ=nZ,ind.drift=ind.drift))
   }
  
  out1<- LatticeKrig.coef( Mc, wPHI, wT.matrix, wy, lambda)
  
  fitted.values<- (wT.matrix%*%out1$d.coef + wPHI%*%out1$c.coef)/sqrt(weights)
  residuals<- y- fitted.values

  out2<-LatticeKrig.lnPlike(Mc, Q, y, N1, N2,lambda,residuals, weights)
# save seed if random number generation happening outside LatticeKrig  
  save.seed<- .Random.seed
  if( !is.na(iseed)){
    set.seed(iseed)}
# generate N(0,1)  
  wEy<- matrix( rnorm( NtrA*N), N,NtrA)*sqrt(weights)
# restore seed   
   if( !is.na(iseed)){
      assign(".Random.seed",save.seed, pos=1)}
#  
  out3<- LatticeKrig.coef( Mc, wPHI, wT.matrix, wEy, lambda)
  wEyhat<-(wT.matrix%*%out3$d.coef + wPHI%*%out3$c.coef) 
  trA.info <- t((wEy *wEyhat)/weights)%*%rep(1,N)
  trA.est <- mean(trA.info)

# find the GCV function
  GCV=  (sum(weights*(residuals)^2)/N )  /( 1- trA.est/np)^2

  
  obj<-list(x=x,y=y,weights=weights, Z=Z,
              d.coef=out1$d.coef, c.coef=out1$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= out2$lnProfileLike,
              rho.MLE=out2$rho.MLE, shat.MLE= out2$shat.MLE,lambda=lambda,
              lnDetCov= out2$lnDetCov, quad.form= out2$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=2,nt=nt, nZ=nZ,ind.drift=ind.drift,
              call=match.call())
  class(obj)<- "LatticeKrig"
  return(obj)
}

LatticeKrig.coef<- function(Mc,wPHI,wT.matrix,wy,lambda){
  A <- forwardsolve(Mc, transpose = TRUE, t(wPHI)%*%wT.matrix, upper.tri = TRUE)
  A <- backsolve(Mc, A)  
  A<- t(wT.matrix)%*%(wT.matrix - wPHI%*%A)/lambda
#   A is  (T^t M^{-1} T) 
  b <- forwardsolve(Mc, transpose = TRUE, t(wPHI)%*%wy, upper.tri = TRUE)
  b <- backsolve(Mc, b)  
  b<- t(wT.matrix)%*%(wy- wPHI%*%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(wPHI)%*%(wy-wT.matrix%*%d.coef),
                         upper.tri = TRUE)
  c.coef <- backsolve(Mc, c.coef)

  return(list( c.coef=c.coef, d.coef= d.coef))  
  }

LatticeKrig.lnPlike<- function(Mc,Q,y,N1,N2,lambda,residuals,weights){
  N<- length(y)
  c.mKrig<- weights*residuals/lambda
# 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(Q))))
# 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)
# 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)
  return( list(lnProfileLike=lnProfileLike,rho.MLE=rho.MLE, shat.MLE=shat.MLE,
               quad.form=quad.form, lnDetCov=lnDetCov) )
 
}


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

predict.LatticeKrig<- function( object, xnew=NULL,Z=NULL,drop.Z=FALSE,...){
  if( is.null(xnew)){
    xnew<- object$x}
  if( is.null(Z)& object$nZ>0){
    Z<- object$Z} 
  PHIg<-  wendland.basis( xnew, object$centers, theta=object$basis.delta,
                          spam.format=TRUE, flavor=0)
  NG<- nrow( xnew)
   if( drop.Z|object$nZ==0){
     temp1<-cbind( rep(1,NG), xnew)%*% object$d.coef[object$ind.drift,]}
   else{
     temp1<- cbind( rep(1,NG), xnew,Z)%*% object$d.coef}
  temp2<- PHIg%*%object$c.coef
return( temp1 + temp2)
}

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( x$nZ>0){
     c1 <- c(c1, "Number of covariates")
     c2 <- c(c2, x$nZ)}
    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$np, " with spacing", x$delta,
        "overlap of ", x$overlap, 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]
  if( spam.format){
# Convert matrix to spam format  
  sM <- spind2spam( list( ind=cbind(Hi,Hj), ra=ra, da=da) )
# convert to an ordinary matrix (usually for debugging)     
  return(sM)}
  else{
    return( list( ind=cbind(Hi,Hj), ra=ra, da=da))}
}

`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.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)))
    }

 
}

LatticeKrig.sim<- function( x1,center.grid.list, beta=-.2,k=3,
                                          overlap=2.5,M=1){
  N1<- length( center.grid.list$x)
  N2<- length( center.grid.list$y)
# total number of basis functions/coefficients
  N<- N1*N2 
  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
#basis fucntions evaluated at points for simulation  
  PHI1<-  wendland.basis( x1, centers, theta=basis.delta, spam.format=FALSE,k=k,
                         flavor=0)
  H<-  lattice.precision(N1,N2, spam.format=TRUE, beta=beta)
  Q<- t(H)%*%(H) 
#  Note:  Q  <- t(H)%*%(H) is the precision matrix of coefficients
#  if Sigma = solve(Q) then to simulate from Sigma need
#     sqrt(Sigma) * N(0,1)   or  solve( sqrt(Q)) *N(0,1)
#
   E<- matrix(rnorm(N*M),N,M)
   chol(Q)-> Mc
# solve linear system   
   A <- forwardsolve(Mc, transpose = TRUE, E, upper.tri = TRUE)
        return( PHI1%*%A)
 
}