## *************************************************************************
##  MLPKriging S3 class for rlibkriging
##
##  Deep Kernel Learning: Kriging with a joint MLP feature extractor
##  Phi : R^d -> R^{d_out}, k(x, x') = sigma^2 * k_base(Phi(x), Phi(x'); theta)
## *************************************************************************

## Register MLPKriging as a known S3 class to the S4 system so that
## packages that override generics (e.g. RobustGaSP overriding simulate)
## can still dispatch to the S3 simulate.MLPKriging method.
setOldClass("MLPKriging")

#' Shortcut to provide functions to the S3 class "MLPKriging"
#' @param obj A list with a \code{ptr} element pointing to a C++ MLPKriging object
#' @return An object of class "MLPKriging" with methods accessible via \code{$}
classMLPKriging <- function(obj) {
    class(obj) <- "MLPKriging"
    for (f in c('as.list','copy','fit','save',
                'logLikelihood','logLikelihoodFun',
                'predict','print','show','simulate','update','update_simulate',
                'is_fitted')) {
        eval(parse(text=paste0("obj$", f, " <- function(...) ", f, "(obj,...)")))
    }
    for (d in c('kernel','theta','sigma2','hidden_dims','activation','feature_dim',
                'X','centerX','scaleX','y','centerY','scaleY',
                'normalize','regmodel','F_','T_','M','z','beta')) {
        eval(parse(text=paste0("obj$", d, " <- function() ", d, "(obj)")))
    }
    obj
}

#' @title Create an MLPKriging model (Deep Kernel Learning)
#'
#' @description Kriging with a joint multi-layer perceptron applied to all
#'   inputs before the GP kernel is evaluated. The MLP weights, GP range
#'   parameters, variance and trend are jointly fitted by maximising the
#'   concentrated log-likelihood.
#'
#' @param y numeric vector of observations (n)
#' @param X numeric matrix of inputs (n x d)
#' @param hidden_dims integer vector of hidden layer sizes, e.g. \code{c(32, 16)}
#' @param d_out output feature dimensionality (default 2)
#' @param activation activation function: "relu", "selu", "tanh", "sigmoid", "elu"
#' @param kernel covariance kernel: "gauss", "matern3_2", "matern5_2", "exp"
#' @param regmodel trend: "constant", "linear", "quadratic"
#' @param normalize logical; normalise inputs?
#' @param optim optimiser (default "BFGS+Adam")
#' @param objective "LL" (log-likelihood)
#' @param parameters optional named list of tuning parameters, e.g.
#'   \code{list(max_iter_adam = "300", adam_lr = "0.001", max_iter_bfgs = "50")}
#'
#' @return An S3 object of class "MLPKriging".
#'
#' @examples
#' X <- as.matrix(seq(0.01, 0.99, length.out = 10))
#' f <- function(x) 1 - 1/2 * (sin(12*x)/(1+x) + 2*cos(7*x)*x^5 + 0.7)
#' y <- f(X)
#' k <- MLPKriging(y, X, hidden_dims = c(16, 8), d_out = 2,
#'                 activation = "selu", kernel = "gauss")
#' print(k)
#'
#' @export
MLPKriging <- function(y, X, hidden_dims,
                       d_out = 2,
                       activation = "selu",
                       kernel = "gauss",
                       regmodel = "constant",
                       normalize = FALSE,
                       optim = "BFGS+Adam",
                       objective = "LL",
                       parameters = NULL) {
  y <- as.numeric(y)
  X <- as.matrix(X)
  hidden_dims <- as.integer(hidden_dims)

  ptr <- mlpKriging_new(y, X, hidden_dims, as.integer(d_out),
                        activation, kernel, regmodel, normalize,
                        optim, objective, parameters)
  obj <- list(ptr = ptr)
  return(classMLPKriging(obj))
}

# -----------------------------------------------------------------------
#  S3 methods
# -----------------------------------------------------------------------

#' @method print MLPKriging
#' @export
print.MLPKriging <- function(x, ...) {
  cat(mlpKriging_summary(x$ptr))
  invisible(x)
}

#' @method summary MLPKriging
#' @export
summary.MLPKriging <- function(object, ...) {
  mlpKriging_summary(object$ptr)
}

#' @title Fit an MLPKriging model to data
#'
#' @description (Re-)fit an already-constructed MLPKriging object on new
#'   data.  The MLP architecture and kernel are kept from construction.
#'
#' @param object MLPKriging object
#' @param y numeric vector of observations (n)
#' @param X numeric matrix of inputs (n x d)
#' @param regmodel trend: "constant", "linear", "quadratic"
#' @param normalize logical; normalise inputs?
#' @param optim optimiser
#' @param objective "LL" (log-likelihood)
#' @param parameters optional named list of tuning parameters
#' @param ... ignored
#'
#' @return No return value. MLPKriging object argument is modified.
#'
#' @method fit MLPKriging
#' @export
fit.MLPKriging <- function(object, y, X,
                           regmodel = "constant",
                           normalize = FALSE,
                           optim = "BFGS+Adam",
                           objective = "LL",
                           parameters = NULL, ...) {
  mlpKriging_fit(object$ptr,
                 as.numeric(y), as.matrix(X),
                 regmodel, normalize, optim, objective,
                 parameters)
  invisible(NULL)
}

#' @title Predict with an MLPKriging model
#' @param object MLPKriging object
#' @param x prediction matrix (m x d)
#' @param return_stdev return standard deviations?
#' @param return_cov return full covariance?
#' @param return_deriv return derivatives of mean and stdev wrt x?
#' @param ... ignored
#' @return list with \code{mean}, optionally \code{stdev}, \code{cov},
#'   \code{mean_deriv}, \code{stdev_deriv}
#' @method predict MLPKriging
#' @export
predict.MLPKriging <- function(object, x, return_stdev = TRUE, return_cov = FALSE,
                               return_deriv = FALSE, ...) {
  mlpKriging_predict(object$ptr, as.matrix(x), return_stdev, return_cov, return_deriv)
}

#' @title Simulate from an MLPKriging model
#' @param object MLPKriging object
#' @param nsim number of simulations
#' @param seed random seed
#' @param x simulation matrix (m x d)
#' @param will_update logical; if \code{TRUE} keep the internal state for a
#'   subsequent \code{update_simulate()} call.
#' @param ... ignored
#' @return matrix (m x nsim)
#' @method simulate MLPKriging
#' @export
simulate.MLPKriging <- function(object, nsim = 1, seed = 123, x,
                                 will_update = FALSE, ...) {
  mlpKriging_simulate(object$ptr, as.integer(nsim),
                      as.integer(seed), as.matrix(x),
                      as.logical(will_update))
}

#' @title Update simulated paths with new observations (FOXY algorithm)
#' @param object MLPKriging object (must have called simulate with will_update=TRUE)
#' @param y_u new observations
#' @param X_u new input matrix
#' @param ... ignored
#' @return matrix (m x nsim) of updated simulated paths
#' @method update_simulate MLPKriging
#' @export
update_simulate.MLPKriging <- function(object, y_u, X_u, ...) {
  mlpKriging_update_simulate(object$ptr, as.numeric(y_u), as.matrix(X_u))
}

#' @title Update an MLPKriging model with new observations
#' @param object MLPKriging object
#' @param y_u new observations
#' @param X_u new input matrix
#' @param refit Logical. If \code{TRUE} the model is refitted (default is TRUE).
#' @param ... ignored
#' @method update MLPKriging
#' @export
update.MLPKriging <- function(object, y_u, X_u, refit = TRUE, ...) {
  mlpKriging_update(object$ptr, as.numeric(y_u), as.matrix(X_u), as.logical(refit))
  invisible(object)
}

#' @method logLikelihood MLPKriging
#' @export
logLikelihood.MLPKriging <- function(object, ...) {
  mlpKriging_logLikelihood(object$ptr)
}

#' @title Evaluate log-likelihood at given GP theta
#' @param object MLPKriging object
#' @param theta range parameter vector
#' @param return_grad return gradient?
#' @param return_hess return hessian?
#' @param ... ignored
#' @method logLikelihoodFun MLPKriging
#' @export
logLikelihoodFun.MLPKriging <- function(object, theta, return_grad = FALSE, return_hess = FALSE, ...) {
  mlpKriging_logLikelihoodFun(object$ptr, theta, return_grad, return_hess)
}

#' @method theta MLPKriging
#' @export
theta.MLPKriging <- function(object, ...) {
  mlpKriging_theta(object$ptr)
}

#' @method sigma2 MLPKriging
#' @export
sigma2.MLPKriging <- function(object, ...) {
  mlpKriging_sigma2(object$ptr)
}

#' @method kernel MLPKriging
#' @export
kernel.MLPKriging <- function(object, ...) {
  mlpKriging_kernel(object$ptr)
}

#' @title Get feature dimensionality (d_out)
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @export
feature_dim <- function(object, ...) UseMethod("feature_dim")

#' @title Get feature dimensionality for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method feature_dim MLPKriging
#' @export
feature_dim.MLPKriging <- function(object, ...) {
  mlpKriging_featureDim(object$ptr)
}

#' @title Get hidden layer sizes
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @export
hidden_dims <- function(object, ...) UseMethod("hidden_dims")

#' @title Get hidden layer sizes for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method hidden_dims MLPKriging
#' @export
hidden_dims.MLPKriging <- function(object, ...) {
  mlpKriging_hiddenDims(object$ptr)
}

#' @title Get activation function name
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @export
activation <- function(object, ...) UseMethod("activation")

#' @title Get activation function for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method activation MLPKriging
#' @export
activation.MLPKriging <- function(object, ...) {
  mlpKriging_activation(object$ptr)
}

#' @title Check whether an MLPKriging model is fitted
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method is_fitted MLPKriging
#' @export
is_fitted.MLPKriging <- function(object, ...) {
  mlpKriging_isFitted(object$ptr)
}

#' @title Get training input matrix
#' @param object MLPKriging object
#' @param ... ignored
#' @return matrix of training inputs
#' @method X MLPKriging
#' @export
X.MLPKriging <- function(object, ...) {
  mlpKriging_X(object$ptr)
}

#' @title Get training output vector
#' @param object MLPKriging object
#' @param ... ignored
#' @return vector of training outputs
#' @method y MLPKriging
#' @export
y.MLPKriging <- function(object, ...) {
  mlpKriging_y(object$ptr)
}

#' @title Get input centering vector for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method centerX MLPKriging
#' @export
centerX.MLPKriging <- function(object, ...) {
  mlpKriging_centerX(object$ptr)
}

#' @title Get input scaling vector for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method scaleX MLPKriging
#' @export
scaleX.MLPKriging <- function(object, ...) {
  mlpKriging_scaleX(object$ptr)
}

#' @title Get output centering value for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method centerY MLPKriging
#' @export
centerY.MLPKriging <- function(object, ...) {
  mlpKriging_centerY(object$ptr)
}

#' @title Get output scaling value for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method scaleY MLPKriging
#' @export
scaleY.MLPKriging <- function(object, ...) {
  mlpKriging_scaleY(object$ptr)
}

#' @title Get normalize flag for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method normalize MLPKriging
#' @export
normalize.MLPKriging <- function(object, ...) {
  mlpKriging_normalize(object$ptr)
}

#' @title Get regression model type for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method regmodel MLPKriging
#' @export
regmodel.MLPKriging <- function(object, ...) {
  mlpKriging_regmodel(object$ptr)
}

#' @title Get trend matrix F for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method F_ MLPKriging
#' @export
F_.MLPKriging <- function(object, ...) {
  mlpKriging_F(object$ptr)
}

#' @title Get Cholesky factor T for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method T_ MLPKriging
#' @export
T_.MLPKriging <- function(object, ...) {
  mlpKriging_T(object$ptr)
}

#' @title Get whitened trend matrix M for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method M MLPKriging
#' @export
M.MLPKriging <- function(object, ...) {
  mlpKriging_M(object$ptr)
}

#' @title Get whitened residuals z for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method z MLPKriging
#' @export
z.MLPKriging <- function(object, ...) {
  mlpKriging_z(object$ptr)
}

#' @title Get trend coefficients beta for an MLPKriging model
#' @param object A Kriging/MLPKriging/WarpKriging model object.
#' @param ... Unused.
#' @method beta MLPKriging
#' @export
beta.MLPKriging <- function(object, ...) {
  mlpKriging_beta(object$ptr)
}

#' @title Deep copy of MLPKriging model
#' @param object MLPKriging object
#' @param ... ignored
#' @return a new independent MLPKriging object
#' @method copy MLPKriging
#' @export
copy.MLPKriging <- function(object, ...) {
  ptr_copy <- mlpKriging_copy(object$ptr)
  classMLPKriging(list(ptr = ptr_copy))
}

#' @title Save an MLPKriging model to file
#' @param object MLPKriging object
#' @param filename path to save file
#' @param ... ignored
#' @method save MLPKriging
#' @export
save.MLPKriging <- function(object, filename, ...) {
  if (!is.character(filename))
    stop("'filename' must be a string")
  mlpKriging_save(object$ptr, filename)
  invisible(NULL)
}

#' @title Load an MLPKriging model from file
#' @param filename path to saved file
#' @param ... ignored
#' @return MLPKriging object
#' @method load MLPKriging
#' @export
load.MLPKriging <- function(filename, ...) {
  if (!is.character(filename))
    stop("'filename' must be a string")
  ptr <- mlpkriging_load(filename)
  classMLPKriging(list(ptr = ptr))
}
