require(mgcv)


#
#  A function to sample equally from the original data range
#
sampleData <- function(model, covariates, newdata, N) {
  for(covariate in covariates) {
    newdata[[covariate]] <- seq(model$var.summary[[covariate]][1], 
                                model$var.summary[[covariate]][3], 
                                length.out=N);
  }
  return(expand.grid(newdata));
}


#
# A "nice" function to evaluate a 2D term of a GAM on its marginals.
#
projectMarginals <- function(model, term_name, new_data, N_samples) {
  #
  # Search for term in model$sp
  #
  term_idx = match(term_name, names(model$sp))
  if (is.na(term_idx)) {
    stop(sprintf("Term %s not found in names(model$sp)!", term_name));
  }
  
  #
  # Extract names of covariates:
  #
  dependent_covariates <- model$smooth[[term_idx]]$term
  if (2 != length(dependent_covariates)) {
    stop(sprintf("Expected exactly two dependent covariates for given smooth, got %i",
                 length(dependent_covariates)));
  }
  x_name = dependent_covariates[[1]];
  y_name = dependent_covariates[[2]];
  
  #
  # Get value range of covariates:
  #
  cov_min = min(c(model$var.summary[[x_name]][[1]],
                  model$var.summary[[y_name]][[1]]));
  cov_max = max(c(model$var.summary[[x_name]][[3]],
                  model$var.summary[[y_name]][[3]]));
  
  #
  # Sample interval:
  #
  cov_samples = seq(cov_min, cov_max, length.out=N_samples)
  
  #
  # Now, create table, where all covariates of model are 0 execpt x_name and y_name:
  #
  new_data[[x_name]] <- cov_samples;
  new_data[[y_name]] <- cov_samples;
  samples <- expand.grid(new_data);
  
  #
  # Get linear-prediction matrix and coefficients of model:
  #
  Xp_full <- predict.gam(model, samples, type="lpmatrix");
  C  <- coef(model);
  
  #
  # Select only columns of Xp that correspond to selected term:
  #
  from_idx <- model$smooth[[term_idx]]$first.para;
  to_idx   <- model$smooth[[term_idx]]$last.para;
  Xp <- matrix(0, dim(Xp_full)[[1]], dim(Xp_full)[[2]])
  Xp[, from_idx:to_idx] <- Xp_full[,from_idx:to_idx];
  rm(Xp_full)
  
  #
  # Perform pediction (mean_xy)
  #
  mean_xy <- Xp %*% C;
  intercept = mean(mean_xy)
  mean_xy = mean_xy - intercept
  
  # Construct Sx (projection matrix on X):
  Sx <- matrix(0.0, N_samples, N_samples); Sx[row(Sx)==col(Sx)] <- 1.0;
  dim(Sx) <- c(N_samples*N_samples); Sx <- rep(Sx, N_samples)
  dim(Sx) <- c(N_samples, N_samples*N_samples);
  
  # Construct Sy (projection)
  Sy <- matrix(0.0, N_samples, N_samples); Sy[row(Sy)==col(Sy)] <- 1.0;
  dim(Sy) <- c(N_samples*N_samples); Sy <- rep(Sy, N_samples)
  dim(Sy) <- c(N_samples*N_samples, N_samples); Sy <- t(Sy)
  dim(Sy) <- c(N_samples*N_samples, N_samples); Sy <- t(Sy)
  
  Px <- Sx %*% mean_xy/N_samples; Vx <- model$sig2 * Sx %*% Xp %*% model$Vp %*% t(Xp) %*% t(Sx)/(N_samples**2)
  Py <- Sy %*% mean_xy/N_samples; Vy <- model$sig2 * Sy %*% Xp %*% model$Vp %*% t(Xp) %*% t(Sy)/(N_samples**2)

  Pxy <- mean_xy - t(Sx) %*% Px - t(Sy) %*% Py
  Vxy <- model$sig2 * Xp %*% model$Vp %*% t(Xp) 
  
  return(list(samples=cov_samples, 
              Px=Px, Vx=Vx, 
              Py=Py, Vy=Vy,
              Pxy=Pxy, Vxy=Vxy,
              Xp=Xp,
              x_name=x_name, y_name=y_name, 
              intercept=intercept));
}