# ---------------------------------------------------
# My own functions for anchor regression:
# ---------------------------------------------------

# Sebastian Sippel
# 08.12.2019
require(matrixStats)
require(weights)


## Function to repeat vector either by row or column:
rep.row<-function(x,n){
  matrix(rep(x,each=n),nrow=n)
}
rep.col<-function(x,n){
  matrix(rep(x,each=n), ncol=n, byrow=TRUE)
}


# Error metrics / weighted root mean squared error:
# https://stats.stackexchange.com/questions/230517/weighted-root-mean-square-error
weighted.rmse <- function(actual, predicted, weight){
  sqrt(sum((predicted-actual)^2*weight)/sum(weight))
}



## Fit ridge regression for given \lambda value:
# no centering or scaling performed
fit.ridge <- function(X, Y, lambda, method = "lm.ridge") {
  if (method == "naive") {
    beta_ct = solve(t(X) %*% X + lambda * diag(p)) %*% t(X) %*% Y
    return(beta_ct)
  } else if (method == "naive2") {
    
    beta_ct = solve(t(X) %*% X + lambda * diag(p), t(X) %*% Y)
    
    return(beta_ct)
  } else if (method == "lm.ridge") {
    return(fit.ridge_s2(X=X, Y=Y, lambda = lambda) )
  }
}



## Fit ridge regression via SVD (following lm.ridge):
# no centering or scaling performed.
fit.ridge_s2 <- function (X, Y, lambda = 0) 
{
  # m <- match.call(expand.dots = FALSE)
  # m$model <- m$x <- m$y <- m$contrasts <- m$... <- m$lambda <- NULL
  # m[[1L]] <- quote(stats::model.frame)
  # m <- eval.parent(m)
  # Terms <- attr(m, "terms")
  # Y <- model.response(m)
  # X <- model.matrix(Terms, m, contrasts)
  n <- nrow(X)
  p <- ncol(X)
  # offset <- model.offset(m)
  # if (!is.null(offset)) 
  #  Y <- Y - offset
  # if (Inter <- attr(Terms, "intercept")) {
  Xm <- colMeans(X)
  Ym <- mean(Y)
  Inter = 0
  # p <- p - 1
  # X <- X - rep(Xm, rep(n, p))
  # Y <- Y - Ym
  # }
  # else Ym <- Xm <- NA
  
  # Xscale <- drop(rep(1/n, n) %*% X^2)^0.5
  # X <- X/rep(Xscale, rep(n, p))
  Xs <- svd(X)
  rhs <- t(Xs$u) %*% Y
  d <- Xs$d
  lscoef <- Xs$v %*% (rhs/d)
  lsfit <- X %*% lscoef
  resid <- Y - lsfit
  s2 <- sum(resid^2)/(n - p - 1)
  HKB <- (p - 2) * s2/sum(lscoef^2)
  LW <- (p - 2) * s2 * n/sum(lsfit^2)
  k <- length(lambda)
  dx <- length(d)
  div <- d^2 + rep(lambda, rep(dx, k))
  a <- drop(d * rhs)/div
  dim(a) <- c(dx, k)
  coef <- Xs$v %*% a
  # dimnames(coef) <- list(names(Xscale), format(lambda))
  GCV <- colSums((Y - X %*% coef)^2)/(n - colSums(matrix(d^2/div, 
                                                         dx)))^2
  res <- list(coef = drop(coef), Inter = Inter, 
              lambda = lambda, ym = Ym, xm = Xm, GCV = GCV, kHKB = HKB, 
              kLW = LW)
  class(res) <- "ridgelm"
  return(res)
}



## Fit anchor regression for given \lambda and \gamma value:
# following Rothenhäusler, D., Meinshausen, N., Bühlmann, P. and Peters, J., 2021. 
# Anchor regression: Heterogeneous data meet causality. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 83(2), pp.215-246.
# fits anchor regression directly
fit.anchor_s <- function(X_sc, Y_sc, A, lambda = 100, gamma = 5, ret.Loss = F) {
  
  X_ct = X_sc
  Y_ct = Y_sc
  
  n = length(A)
  p = dim(X_ct)[2]
  
  PA = A %*% solve(t(A) %*% A) %*% t(A)
  PAc = diag(n) - PA
  
  D_l = lambda * diag(p) + gamma * t(X_ct) %*% PA %*% X_ct + t(X_ct) %*% PAc %*% X_ct
  d = (-gamma * t(PA %*% Y_ct) %*% X_ct - t(PAc %*% Y_ct) %*% X_ct) / (-1) 
  
  beta_lg = solve(D_l) %*% t(d)
  
  if (ret.Loss == F) {
    return(c(beta_lg))
  } else if (ret.loss == T) {
    ret.list = list()
    ret.list$beta = beta_lg
    ret.list$PA = PA  
    ret.list$Loss = 
      c(perf = sum(c((diag(n) - PA) %*% (Y_ct - X_ct %*% beta_lg))^2),
      anchor = gamma * sum(( PA %*% (Y_ct - X_ct %*% beta_lg) )^2),
      ridge = lambda * sqrt(sum(beta_lg^2)))
        # sum(perf + anchor + ridge)
    return(ret.list$Loss)
  }
}






## Fit anchor regression for given \lambda and \gamma value:
# following Rothenhäusler, D., Meinshausen, N., Bühlmann, P. and Peters, J., 2021. 
# Anchor regression: Heterogeneous data meet causality. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 83(2), pp.215-246.
# fits anchor regression via a transformation (as described in Rothenhäusler et al., 2021) and subsequent ridge regression.
# standardization optional. Default is to not standardize.
fit.anchor_t <- function(y, x, A, lambda = 100, gamma = 10, standardize = F) {
  if (standardize == T) {
    # Standardize x and y before anchor regression:
    mu_x=colMeans(x); sd_x=colSds(x); # image.plot(matrix(sd_x, 72, 36))
    mu_y=mean(y); sd_y=sd(y)
    X_sc = (x - rep.row(mu_x, n = length(y))) / rep.row(sd_x, n = length(y))
    Y_sc = (y - mu_y) / sd_y
    
    anchor = anchor.transform.data(X_sc = X_sc, Y_sc = Y_sc, A = A, gamma = gamma)
    beta = fit.ridge_s2(X = anchor$X_tr, Y = anchor$Y_tr, lambda = lambda)$coef
    
    # convert to un-standardized coefficients:
    # https://stats.stackexchange.com/questions/155362/glmnet-unstandardizing-linear-regression-coefficients
    require(matrixStats)
    a0 = mu_y - colSums(beta * mu_x / sd_x) * sd_y
    beta0 = beta * sd_y / sd_x
    return(list(a0 = a0, beta = beta0, lambda = lambda, gamma = gamma, beta_std = beta))
  } else if (standardize == F) {
    anchor = anchor.transform.data(X_sc = X_sc, Y_sc = Y_sc, A = A, gamma = gamma)
    beta = fit.ridge_s2(X = anchor$X_tr, Y = anchor$Y_tr, lambda = lambda)$coef
    return(beta)
  }
}



## Function to implement anchor regression transformation to subsequently perform ridge regression.
anchor.transform.data <- function(X_sc, Y_sc, A, gamma = 10, PA=NULL) {
  
  n = length(Y_sc)
  p = dim(X_sc)[2]
  
  if (is.null(PA)) PA <- A %*% solve(t(A)%*%A) %*% t(A)
  TR <- ( sqrt(gamma)*PA + (diag(n)-PA))
  
  X_tr = TR %*% X_sc
  Y_tr = TR %*% Y_sc
  return(list(X_tr = X_tr, Y_tr = c(Y_tr), A = A))
}





## Fit anchor regression with cross-validation.
# Function implements cross-validation and calls fit.anchor_t.
# Calculations in parallel over number of simulations (with cv = "model.subagging") or number of models (with cv = "leave.model.out")
# Function calculates several error metrics.
cv.anchor <- function(x, y, A, lambda, gamma, foldid, 
                      nr.cores = NULL, nr.subsample = 3000, cv = "model.subagging", nsim = 20, keep = F, adj.mean.by.mod = F) {
  
  # Purpose of function: 
  # Similar to cv.glmnet, but including options for different \gamma values
  ## Standardization of regression coefficients is based on: https://stats.stackexchange.com/questions/155362/glmnet-unstandardizing-linear-regression-coefficients
  
  # 1. Prepare regression/cross-validation:
  foldid.un = na.omit(unique(foldid))

  # Different cross-validation strategy to use:
  if (cv == "model.subagging") {
    nuse=0.5
    design=sapply(X = 1:nsim, FUN=function(ix) {
      set.seed(ix+6) 
      sample(x = 1:length(foldid.un), size = ceiling(length(foldid.un)*nuse), replace = F) })
  } else if (cv == "leave.model.out") {
    nsim = length(foldid.un)
    design = sapply(X = 1:nsim, FUN=function(ix) c(1:nsim)[-ix])
  } else if (cv == "model.by.model") {     # -> need to set nr.subsample == NULL
    # nsim = length(um)
    # design <- matrix(data = c(1:nsim), nrow = 1, ncol = nsim)
  }
  
  require(doParallel)
  registerDoParallel(cores = nr.cores)
  # ptm <- proc.time()
  anchor.list = foreach(sim=1:nsim) %dopar% {
    
    print(paste("\r ***", sim))
    train <- numeric(0)
    for (ucc in 1:length(foldid.un)){ if( ucc %in% design[,sim]) train <- c(train, foldid.un[ucc])}
    test <- foldid.un[-which(foldid.un %in% train)]
    
    itrainX <- which(foldid %in% train)
    itestX <- which(!(foldid %in% train))
    
    ## Subsample training indices and standardize data:
    # Subsample itrainX to a train ix of equal weight:
    set.seed(sim)
    
    if (is.null(nr.subsample)) {
      train.ix = itrainX
    } else {
      train.ix = c(sapply(X = design[,sim], FUN=function(cc) sort(sample(x = which(cc == foldid), size = nr.subsample, replace=T))))
    }
    cv.anchor.out = list()
    cv.anchor.out$anchor.fit = fit.anchor_t(x = x[train.ix,], y = y[train.ix], A = A[train.ix], lambda = lambda, gamma = gamma, standardize = T)
    cv.anchor.out$Yhat = matrix(data = NA, nrow = length(y), ncol = length(lambda))
    if (adj.mean.by.mod == F) {
      cv.anchor.out$Yhat[itestX,] = x[itestX,] %*% cv.anchor.out$anchor.fit$beta + rep.row(cv.anchor.out$anchor.fit$a0, n = length(itestX))
    } else if (adj.mean.by.mod == T) {
      pred = x[itestX,] %*% cv.anchor.out$anchor.fit$beta # + rep.row(cv.anchor.out$anchor.fit$a0, n = length(itestX))
      cv.anchor.out$Yhat[itestX,] = pred - rep.row(colMeans(pred), length(itestX)) + mean(y[itestX])
    }
    return(cv.anchor.out)
  }
  # proc.time() - ptm
  
  ## SUMMARY STATISTICS:
  # Weighted RMSE + Weighted correlation
  Yhat = sapply(X = 1:length(lambda), FUN=function(lambda.ix) rowMeans(sapply(X = anchor.list, FUN=function(x) x$Yhat[,lambda.ix]), na.rm=T))
  Yhat.sd = sapply(X = 1:length(lambda), FUN=function(lambda.ix) rowSds(sapply(X = anchor.list, FUN=function(x) x$Yhat[,lambda.ix]), na.rm=T))
  beta = sapply(X = 1:length(lambda), FUN=function(lambda.ix) rowMeans(sapply(X = anchor.list, FUN=function(x) x$anchor.fit$beta[,lambda.ix]), na.rm=T))
  beta_std = sapply(X = 1:length(lambda), FUN=function(lambda.ix) rowMeans(sapply(X = anchor.list, FUN=function(x) x$anchor.fit$beta_std[,lambda.ix]), na.rm=T))
  a0 = sapply(X = 1:length(lambda), FUN=function(lambda.ix) mean(sapply(X = anchor.list, FUN=function(x) x$anchor.fit$a0[lambda.ix]), na.rm=T))
  
  # Save prediction from each simulation for later lambda selection:
  allsim = list()
  allsim$design = design
  allsim$beta = lapply(X = anchor.list, FUN=function(x) x$anchor.fit$beta)
  allsim$a0 = lapply(X = anchor.list, FUN=function(x) x$anchor.fit$a0)
  if (keep == T) allsim$Yhat = lapply(X = anchor.list, FUN=function(x) x$Yhat)
    
  w=rep(NA, length(y)); for (cc in 1:length(foldid.un)) w[which(cc == foldid)] = 1/length(which(cc == foldid)) / length(foldid.un)
  MSE = sapply(X = 1:length(lambda), FUN=function(lambda.ix) weighted.rmse(actual = y, predicted = Yhat[,lambda.ix], weight = w)^2)
  cur.resid = Yhat - rep.col(y, n = length(lambda))
  # res.cor = sapply(X = 1:length(lambda), FUN=function(lambda.ix) wtd.cor(x = cur.resid[,lambda.ix], y = A, weight = w)[1])
  
  ## MSE and residual correlation by model:
  MSE.by.mod = matrix(data = NA, nrow = length(foldid.un), ncol = length(lambda))
  MSE.by.mod[foldid.un,] = t(sapply(foldid.un, FUN=function(cc) {
    c.ix = which(cc == foldid)
    sapply(X = 1:length(lambda), FUN=function(lambda.ix) mse(sim = Yhat[c.ix,lambda.ix], obs = y[c.ix]))
  }))
  # MSE standard deviation across models:
  ## CONTINUE HERE: STANDARD ERROR OF THE MEAN (i.e. variation of the MEAN given folds...)
  SE.MSE =  colSds(x = MSE.by.mod) / sqrt(length(foldid.un))
  #plot(MSE)
  #lines(MSE + SE.MSE, col="red")  # -> 1SE criterion...
  cor.by.mod = matrix(data = NA, nrow = length(foldid.un), ncol = length(lambda))
  cor.by.mod[foldid.un,] = t(sapply(foldid.un, FUN=function(cc) {
    c.ix = which(cc == foldid)
    sapply(X = 1:length(lambda), FUN=function(lambda.ix) cor(x = Yhat[c.ix,lambda.ix], y = y[c.ix]))
  }))
  
  res.cor.by.mod = matrix(data = NA, nrow = length(foldid.un), ncol = length(lambda))
  res.cor.by.mod[foldid.un,] = t(sapply(foldid.un, FUN=function(cc) {
    c.ix = which(cc == foldid)
    sapply(X = 1:length(lambda), FUN=function(lambda.ix) cor(x = cur.resid[c.ix,lambda.ix], y = A[c.ix]))
  }))
  
  # Return list:
  ret.list = list()
  ret.list$beta = beta
  ret.list$a0 = a0
  ret.list$beta_std = beta_std
  ret.list$Yhat = Yhat
  ret.list$Yhat.sd = Yhat.sd
  ret.list$Y = y
  ret.list$A = A
  ret.list$foldid = foldid
  ret.list$MSE = MSE
  ret.list$SE.MSE = SE.MSE
  # ret.list$res.cor = res.cor
  ret.list$MSE.by.mod = MSE.by.mod
  ret.list$cor.by.mod = cor.by.mod
  ret.list$res.cor.by.mod = res.cor.by.mod
  ret.list$allsim = allsim
  
  return(ret.list)
}