#
# fields  is a package for analysis of spatial data written for
# the R software environment.
# Copyright (C) 2024 Colorado School of Mines
# 1500 Illinois St., Golden, CO 80401
# Contact: Douglas Nychka,  douglasnychka@gmail.com,
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with the R software environment if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
# or see http://www.r-project.org/Licenses/GPL-2
##END HEADER
"Krig" <- function(x, Y, cov.function = "stationary.cov", 
                   lambda = NA, df = NA, GCV = FALSE, 
                   XMat = NULL,
                   Z=NULL, 
                   cost = 1,
                   weights = NULL, m = 2, nstep.cv = 200, scale.type = "user", 
                   x.center = rep(0, ncol(x)), x.scale = rep(1, ncol(x)), sigma = NA, 
                   tau2 = NA, method = "REML", verbose = FALSE,
                   null.function = "Krig.null.function", wght.function = NULL, 
                   offset = 0,  na.rm = TRUE, cov.args = NULL, 
                   chol.args = NULL, null.args = NULL, wght.args = NULL, W = NULL, 
                   give.warnings = TRUE,
                   mean.obj = NA, 
                   sd.obj = NA,
                   ...)
  # the verbose switch prints many intermediate steps as an aid in debugging.
  #
{ 
  #
  # use of "Z"  for fixed part of model has been switched to XMat  
  if( !is.null( Z)){
    XMat<-Z
    warning(" Z  as  an arguments has been renamed to XMat within the function
    to be more consistent with spatial process model notation. 
    Please use the XMat argument in the call to Krig avoid this warning. ")
  }
  #
  #
  # create output list
  out <- list()
  ###########################################################
  #  First series of steps simply store pieces of the passed
  #    information to the output list (i.e. the Krig object)
  ##########################################################
  out$call <- match.call()
  #   turn off warning based on options
  if( options()$warn < 0 ){
    give.warnings<- FALSE
  }
  if( !is.na(mean.obj)|!is.na(sd.obj)){
    stop( "correlation model no longer supported please use 
          fields version 11.6")
    
  } 
  #
  # save covariance function as its name
  #
  if( !is.character( cov.function)){
    out$cov.function.name <- as.character(substitute(cov.function))
  }
  else{ 
    out$cov.function.name<-cov.function
  } 
  #
  # save null space function as its name
  #
  out$null.function.name <- as.character(substitute(null.function))
  #
  # save weight  function as its name if it is not a NULL
  #
  if (is.null(wght.function)) {
    out$wght.function.name <- NULL
  }
  else {
    out$wght.function.name <- as.character(substitute(wght.function))
  }
  out$W <- W
  if (verbose) {
    print(out$cov.function.name)
    print(out$null.function.name)
    print(out$wght.function.name)
  }
  #
  # logical to indicate if the 'C' argument is present in cov.function
  # -- a bit of esoteric R code!
  C.arg.missing <- all(names(formals(get(out$cov.function.name))) != 
                         "C")
  if (C.arg.missing) 
    stop("Need to have C argument in covariance function\nsee Exp.cov.simple as an example")
  #
  # save parameters values possibly passed to the covariance function
  # also those added to call are assumed to be covariance arguments.
  if (!is.null(cov.args)) 
    out$args <- c(cov.args, list(...))
  else out$args <- list(...)
  #
  # default values for null space function
  out$null.args <- null.args
  #
  #       set degree of polynomial null space if this is default
  #       mkpoly is used so often is it helpful to include m argument
  #       by default in Krig call.
  if (out$null.function.name == "Krig.null.function") {
    out$null.args <- list(m = m)
    out$m <- m
  }
  #
  # default values for Cholesky decomposition, these are important
  # for sparse matrix decompositions used in Krig.engine.fixed.
  if (is.null(chol.args)) {
    out$chol.args <- list(pivot = FALSE)
  }
  else {
    out$chol.args <- chol.args
  }
  # additional arguments for weight matrix.
  out$wght.args <- wght.args
  #
  # the offset is the effective number of parameters used in the GCV
  # calculations -- unless this is part of an additive model this
  # is likely Zero
  out$offset <- offset
  #
  # the cost is the multiplier applied to the GCV eff.df
  # 
  # lambda and df are two ways of parameterizing the smoothness
  # and are related by a monotonic function that unfortunately
  # depends on the locations of the data.
  # lambda can be used directly in the linear algebra, df
  # must be transformed to lambda numerically using the monotonic trransformation
  # tau2 is the error variance and sigma the multiplier for the covariance
  # method is how to determine lambda
  # the GCV logical forces the code to do the more elaborate decompositions
  # that faclitate estimating lambda -- even if a specific lambda value is
  # given.
  out$cost <- cost
  out$lambda <- lambda
  out$eff.df <- df
  out$tau2 <- tau2
  out$sigma <- sigma
  out$method <- method
  out$GCV <- GCV
  #
  # correlation model information
  # set this to FALSE -- has been depreciated.
  #
  out$correlation.model <- FALSE
  #
  # transformation info
  out$scale.type <- scale.type
  out$x.center <- x.center
  out$x.scale <- x.scale
  #
  if (verbose) {
    cat("  Cov function arguments in call  ", fill = TRUE)
    print(out$args)
    cat(" covariance function used is : ", fill = TRUE)
    print(out$cov.function.name)
  }
  ###############################################################
  # Begin modifications and transformations of input information
  # note that many of these manipulations follow a strategy
  # of passing the Krig object (out) to a function and
  # then appending the information from this function to
  # the Krig object (usually also called "out").
  #In this way the Krig object  is built up
  # in steps and the process is easier to follow.
  ###############################################################
  # various checks on x and  Y including removal of NAs in Y
  # Here is an instance of adding to the Krig object
  # in this case also some onerous bookkeeping making sure arguments are consistent
  out2 <- Krig.check.xY(x, Y, XMat, weights, na.rm, verbose = verbose)
  out <- c(out, out2)
  # find replicates and collapse to means and pool variances.
  # Transform unique x locations 
  out2 <- Krig.transform.xY(out,  verbose = verbose)
  
  out <- c(out, out2)
  
  # NOTE: knots have been transformed after this step
  #############################################################
  #  Figure out what to do
  #############################################################
  #
  # this functions works through the logic of
  # what has been supplied for lambda
  out2 <- Krig.which.lambda(out)
  out[names(out2)] <- out2  
  # Make weight matrix for observations
  #    ( this is proportional to the inverse square root of obs covariance)
  #     if a weight function or W has not been passed then this is
  #     diag( out$weightsM) for W
  #     The checks represent a limitation of this model to
  #     the  WBW type decoposition and no replicate observations.
  out$nondiag.W <- (!is.null(wght.function)) | (!is.null(W))
  # Do not continue if there there is a nondiagonal weight matrix
  # and replicate observations.
  if (out$nondiag.W) {
    if (out$knot.model | out$fixed.model) {
      stop("Non diagonal weight matrix for observations
                    not supported\nwith knots or fixed lambda.")
    }
    if (!is.na(out$tauHat.pure.error)) {
      stop("Non diagonal weight matrix not implemented
                    with replicate locations")
    }
  }
  #  make weight matrix and its square root having passed checks
  out <- c(out, Krig.make.W(out, verbose = verbose))
  #   Do the intensive linear algebra to find the solutions
  #   this is where all the heavy lifting happens.
  #
  #   Note that all the information is passed as a list
  #   including arguments to the cholesky decomposition
  #   used within Krig.engine.fixed
  #
  # The results are saved in the component matrices
  #
  # if method=='user' then just evaluate at single lambda
  #  fixed here means a fixed lambda
  #
  # For fixed lambda the decompositions with and without knots
  # are surprisingly similar and so are in one engine.
  ###########################################################
  if (out$fixed.model) {
    out$matrices <- Krig.engine.fixed(out, verbose = verbose)
    #  The trace of A matrix in fixed lambda case is not easily computed
    #  so set this to NA.
    out$eff.df <- NA
  }
  # alternative are
  # matrix decompositions suitable for
  # evaluation at many lambdas to facilitate GCV/REML estimates  etc.
  #
  else{
  # standard engine following the basic computations for thin plate splines
      out$matrices <- Krig.engine.default(out, verbose = verbose)
  # store basic information about decompositions
  }
  out$nt <- out$matrices$nt
  out$np <- out$matrices$np
  out$decomp <- out$matrices$decomp
  #
  # Now determine a logical vector to indicate coefficients tied to  the
  # the 'spatial drift' i.e. the fixed part of the model
  # that is not due to the XMat covariates.
  # NOTE that the spatial drift coefficients must be the first columns of the
  # M matrix
  if (is.null(out$XMat)) {
    out$ind.drift <- rep(TRUE, out$nt)
  }
  else {
    
    mXMat <- ncol(out$XMatM)
    out$ind.drift <- c(rep(TRUE, out$nt - mXMat), rep(FALSE, 
                                                   mXMat))
  }
  if (verbose) {
    cat("null df: ", out$nt, "drift df: ", sum(out$ind.drift), 
        fill = TRUE)
  }
  #################################################
  # Do GCV and REML search over lambda if not fixed or if GCV variable is TRUE
  #  also does a search over likelihood for lambda.
  #################################################
  if (!out$fixed.model | out$GCV) {
    if (verbose) {
      cat("call to KrigFindLambda", fill = TRUE)
    }
    gcv.out <- KrigFindLambda(out, nstep.cv = nstep.cv, verbose = verbose, 
                        cost = out$cost, offset = out$offset, give.warnings=FALSE)
    out$gcv.grid <- gcv.out$gcv.grid
    #   save a handy summary table of the search results
    out$lambda.est <- gcv.out$lambda.est
    
    out$warningTable<- gcv.out$warningTable
    if( verbose){
      cat("summaries from grid search/optimization", fill=TRUE)
      print(out$lambda.est)
      print(out$warningTable)
    }
    if( give.warnings){
      #NOTE: only print out grid search warning for the method of interest.
      printGCVWarnings( gcv.out$warningTable, method=method)
    }
    # assign the preferred lambda either from GCV/REML/MSE or the user value
    # NOTE: gcv/reml can be done but the estimate is
    # still evaluted at the passed user values of lambda (or df)
    # If df is passed need to calculate the implied lambda value
    if (out$method != "user") {
      out$lambda <- gcv.out$lambda.est[out$method, 1]
      out$eff.df <- out$lambda.est[out$method, 2]
    }
    else {
      if (!is.na(out$eff.df)) {
        out$lambda <- Krig.df.to.lambda(out$eff.df, out$matrices$D)
      }
      else {
        out$eff.df <- Krig.ftrace(out$lambda, out$matrices$D)
      }
      # add in  values to GCV table using these values
      # this the 7th row. 
      lam.user <- out$lambda
      info<- gcv.out$info
      newRow<- c( lam.user,      
                  out$eff.df,
                  Krig.fgcv(lam.user, info),
                  sqrt(Krig.fs2hat(lam.user, info)),
                  Krig.flplike(lam.user, info),
                  NA
      )
      out$lambda.est <- 
        rbind( out$lambda.est,user = newRow)
           
    }
  }
  ##########################
  # end GCV/REML block
  ##########################
  #
  # Now we clean up what has happened and stuff 
  # information into output object.
  #
  ##########################################
  # find coeficients at prefered lambda
  # and evaluate the solution at observations
  ##########################################
  #   pass replicate group means -- no need to recalculate these.
  out2 <- Krig.coef(out, yM = out$yM, verbose = verbose)
  out <- c(out, out2)
  #######################################################################
  # fitted values and residuals and predicted values for full model and
  # also on the null space (fixed
  # effects). But be sure to do this at the nonmissing x's.
  ##################################################################
  out$fitted.values <- predict.Krig(out, x = out$x, XMat = out$XMat, 
                                    eval.correlation.model = FALSE)
  out$residuals <- out$y - out$fitted.values
  #
  # this is just M%*%d  note use of do.call using function name
  Tmatrix <- do.call(out$null.function.name, c(out$null.args, 
                                               list(x = out$x, XMat = out$XMat)))
  out$fitted.values.null <- as.matrix(Tmatrix) %*% out$d
  #
  # verbose block
  if (verbose) {
    cat("residuals", out$residuals, fill = TRUE)
  }
  #
  # find various estimates of tau and sigma
  out2 <- Krig.parameters(out)
  out <- c(out, out2)
  ################################################
  # assign the 'best' model as a default choice
  # either use the user supplied values or the results from
  # optimization
  ################################################
  passed.tau2 <- (!is.na(out$tau2))
  if (out$method == "user" & passed.tau2) {
    out$best.model <- c(out$lambda, out$tau2, out$sigma)
  }
  else {
    # in this case lambda is from opt. or supplied by user
    out$best.model <- c(out$lambda, out$tauHat.MLE^2, out$sigmahat)
  }
  # Note: values in best.model are used in subsquent functions as the choice
  # for these parameters!
  # set class
  
##########################
  out$rhohat<- out$sigmahat
  class(out) <- c("Krig")
  return(out)
}
