##################################################################################################################
# author: Jona Cederbaum
##################################################################################################################
##################################################################################################################
# description: function to estimate smooth auto-covariance operators using the triangle only and without symmetry
# constraint. Gives out auto-covariances evaluated on a pre-specified grid.
##################################################################################################################
estimate_cov_tri_constr_fun <- function(index_upperTri, use_RI = FALSE, use_simple = FALSE, bf, method = "REML", 
                                        d_grid = 100, grid_row, grid_col, same_subject_grid, same_word_grid, 
                                        same_curve_grid, same_point_grid, bs, m, use_bam, t,  
                                        mp = TRUE, para_estim = FALSE, para_estim_nc, weights = NULL, np = TRUE,
                                        use_discrete = TRUE){
  
  if(use_bam == TRUE){
    #######################
    # covariance estimation
    # using bam
    #######################
    
    ################
    # set cluster
    # for estimation
    # if specified
    ################
    if(para_estim & !use_discrete){
      if(detectCores() > 1){
        nc_use <- min(detectCores(), para_estim_nc)
        if(.Platform$OS.type=="unix"){
          cl_estim <- makeForkCluster(nnodes = nc_use) # only runs on linux
        }else{
          cl_estim <- makeCluster(nc_use) # also runs on windows
          clusterExport(cl = cl_estim, c("Predict.matrix.symm.smooth", 
                        "make_summation_matrix", "smooth.construct.symm.smooth.spec",
                        "PredictMat", "tensor.prod.model.matrix"))
        }
      }else{
        cl_estim <- NULL
      }
    }else{
      cl_estim <- NULL
    }
    
    if(!use_RI){
      ##################
      # for crossed fRIs
      if(use_discrete){
        time_cov_estim <- 
          system.time(gam1 <- try(bam(cross_vec_bivariate ~ - 1 + 
                                        s(row_t_bivariate, col_t_bivariate, by = same_subject, k = bf[1], 
                                          bs = "symm", m = m[[1]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                        s(row_t_bivariate, col_t_bivariate, by = same_word, k = bf[2], 
                                          bs = "symm", m = m[[2]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                        s(row_t_bivariate, col_t_bivariate, by = same_curve, k = bf[3], 
                                          bs = "symm", m = m[[3]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                        same_point, method = method, data = index_upperTri,  
                                      discrete = TRUE, nthreads = para_estim_nc, weights = weights)))  
      }else{
        time_cov_estim <- 
          system.time(gam1 <- try(bam(cross_vec_bivariate ~ - 1 + 
                                        s(row_t_bivariate, col_t_bivariate, by = same_subject, k = bf[1], 
                                          bs = "symm", m = m[[1]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                        s(row_t_bivariate, col_t_bivariate, by = same_word, k = bf[2], 
                                          bs = "symm", m = m[[2]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                        s(row_t_bivariate, col_t_bivariate, by = same_curve, k = bf[3], 
                                          bs = "symm", m = m[[3]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                        same_point, method = method, data = index_upperTri, 
                                      cluster = cl_estim, weights = weights)))
      }
    }else{       
      if(!use_simple){
        #############
        # for one fRI
        if(use_discrete){
          time_cov_estim <- 
            system.time(gam1 <- try(bam(cross_vec_bivariate ~ 1 + 
                                          s(row_t_bivariate, col_t_bivariate, k = bf[1], bs = "symm", m = m[[1]], 
                                            xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                          s(row_t_bivariate, col_t_bivariate, by = same_curve, k = bf[2], 
                                            bs = "symm", m = m[[2]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                          same_point, method = method, data = index_upperTri, 
                                        discrete = TRUE, nthreads = para_estim_nc, weights = weights)))  
        }else{
          time_cov_estim <- 
            system.time(gam1 <- try(bam(cross_vec_bivariate ~ 1 + 
                                          s(row_t_bivariate, col_t_bivariate, k = bf[1], bs = "symm", m = m[[1]], 
                                            xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                          s(row_t_bivariate, col_t_bivariate, by = same_curve, k = bf[2], 
                                            bs = "symm", m = m[[2]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                          same_point, method = method, data = index_upperTri, cluster = cl_estim, weights = weights)))    
        }
      }else{
        ########################
        # for independent curves
        if(use_discrete){
          time_cov_estim <- 
            system.time(gam1 <- try(bam(cross_vec_bivariate ~ 1 + 
                                          s(row_t_bivariate, col_t_bivariate, k = bf[1], bs = "symm", m = m[[1]], 
                                            xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                          same_point, method = method, data = index_upperTri, 
                                        discrete = TRUE, nthreads = para_estim_nc, weights = weights)))            
        }else{
          time_cov_estim <- 
            system.time(gam1 <- try(bam(cross_vec_bivariate ~ 1 + 
                                          s(row_t_bivariate, col_t_bivariate, k = bf[1], bs = "symm", m = m[[1]], 
                                            xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                          same_point, method = method, data = index_upperTri, cluster = cl_estim, weights = weights)))  
        }
      }
    }
    
    ##########################
    # stop cluster if existing
    ##########################
    if (!is.null(cl_estim)) stopCluster(cl_estim) 
  }else{      
    #######################
    # covariance estimation
    # using gam
    #######################
    
    if(!use_RI){
      ##################
      # for crossed fRIs
      time_cov_estim <- 
        system.time(gam1 <- try(gam(cross_vec_bivariate ~ - 1 + 
                                      s(row_t_bivariate, col_t_bivariate, by = same_subject, k = bf[1], 
                                        bs = "symm", m = m[[1]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                      s(row_t_bivariate, col_t_bivariate, by = same_word, k = bf[2], 
                                        bs = "symm", m = m[[2]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                      s(row_t_bivariate, col_t_bivariate, by = same_curve, k = bf[3], 
                                        bs = "symm", m = m[[3]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                      same_point, method = method, data = index_upperTri, weights = weights)))
    }else{   
      if(!use_simple){
        #############
        # for one fRI
        time_cov_estim <- 
          system.time(gam1 <- try(gam(cross_vec_bivariate ~ 1 + 
                                        s(row_t_bivariate, col_t_bivariate, k = bf[1], bs = "symm", m = m[[1]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                        s(row_t_bivariate, col_t_bivariate, by = same_curve, k = bf[2], 
                                          bs = "symm", m = m[[2]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                        same_point, method = method, data = index_upperTri, weights = weights)))
      }else{
        ########################
        # for independent curves
        time_cov_estim <- 
          system.time(gam1 <- try(gam(cross_vec_bivariate ~ 1 + 
                                        s(row_t_bivariate, col_t_bivariate, k = bf[1], bs = "symm", m = m[[1]], xt = list(bsmargin = "ps", kroneckersum = mp, np = np)) + 
                                        same_point, method = method, data = index_upperTri, weights = weights)))
      }
    }
  }
  
  
  if(class(gam1)[1] != "try-error"){    
    
    ##############################
    # extract smoothing parameters
    ##############################
    sp <- gam1$sp
    
    if(use_RI){
      ###################
      # extract intercept
      ###################
      intercept <- as.numeric(coefficients(gam1)[1])
    }
    
    #################
    # extract sigmasq
    # and truncate
    #################
    sigmasq <- as.numeric(max(coefficients(gam1)["same_point"], 0))
    
    #################################
    # compute the integral of sigmasq
    #################################
    sigmasq_int <- (max(unlist(t))-min(unlist(t))) * sigmasq
    
    ##########################
    # construct data frame for
    # evaluation on grid 
    ##########################
    
    if(!use_RI){
      ##################
      # for crossed fRIs
      grid_data <- data.table(row_t_bivariate = grid_row, col_t_bivariate = grid_col, same_subject = same_subject_grid, 
                              same_word = same_word_grid, same_curve = same_curve_grid, same_point = same_point_grid)
      
      ##########################
      # remove unnecessary stuff
      ##########################
      grid_row <- NULL
      grid_col <- NULL
      same_subject_grid <- NULL
      same_word_grid <- NULL
      same_curve_grid <- NULL
      same_point_grid <- NULL
      
    }else{
      ###################################
      # for one fRI or independent curves
      grid_data <- data.table(row_t_bivariate = grid_row, col_t_bivariate = grid_col, same_subject = same_subject_grid, 
                              same_curve = same_curve_grid, same_point = same_point_grid)
      
      ##########################
      # remove unnecessary stuff
      ##########################
      grid_row <- NULL
      grid_col <- NULL
      same_subject_grid <- NULL
      same_curve_grid <- NULL
      same_point_grid <- NULL
    }
    
    ####################
    # evaluation on grid
    ####################
    time_cov_pred_grid <- 
      system.time(grid_smooth <- predict(gam1, newdata = grid_data[row_t_bivariate <=  col_t_bivariate, ], 
                                         na.omit = TRUE, type = "terms"))
    
    #####################
    # extract covariances
    #####################
    if(!use_RI){
      ##################
      # for crossed fRIs
      grid_mat_B <- matrix(NA, nrow = d_grid, ncol = d_grid)
      grid_mat_B[lower.tri(grid_mat_B, diag = TRUE)] <- grid_smooth[, 2]
      grid_mat_B <- t(grid_mat_B)
      grid_mat_B[lower.tri(grid_mat_B, diag = TRUE)] <- grid_smooth[, 2]
      
      grid_mat_C <- matrix(NA, nrow = d_grid, ncol = d_grid)
      grid_mat_C[lower.tri(grid_mat_C, diag = TRUE)] <- grid_smooth[, 3]
      grid_mat_C <- t(grid_mat_C)
      grid_mat_C[lower.tri(grid_mat_C, diag = TRUE)] <- grid_smooth[, 3]
      
      grid_mat_E <- matrix(NA, nrow = d_grid, ncol = d_grid)
      grid_mat_E[lower.tri(grid_mat_E, diag = TRUE)] <- grid_smooth[, 4]
      grid_mat_E <- t(grid_mat_E)
      grid_mat_E[lower.tri(grid_mat_E, diag = TRUE)] <- grid_smooth[, 4]
      
    }else{
      ###################################
      # for one fRI or independent curves
      grid_mat_B <- matrix(NA, nrow = d_grid, ncol = d_grid)
      grid_mat_B[lower.tri(grid_mat_B, diag = TRUE)] <- grid_smooth[, 2] + intercept
      grid_mat_B <- t(grid_mat_B)
      grid_mat_B[lower.tri(grid_mat_B, diag = TRUE)] <- grid_smooth[, 2] + intercept
      
      grid_mat_C <- matrix(rep(0, length = d_grid^2), ncol = d_grid, nrow = d_grid, byrow = TRUE)  
      
      if(!use_simple){
        #############
        # for one fRI
        grid_mat_E <- matrix(NA, nrow = d_grid, ncol = d_grid)
        grid_mat_E[lower.tri(grid_mat_E, diag = TRUE)] <- grid_smooth[, 3]
        grid_mat_E <- t(grid_mat_E)
        grid_mat_E[lower.tri(grid_mat_E, diag = TRUE)] <- grid_smooth[, 3]  
      }else{
        grid_mat_E <- matrix(rep(0, length = d_grid^2), ncol = d_grid, nrow = d_grid, byrow = TRUE)  
      }
    }
    grid_smooth <- NULL
    grid_data <- NULL
    
  }else{
    sigmasq <- NA
    sigmasq_int <- NA
    grid_mat_B <- NA
    grid_mat_C <- NA
    grid_mat_E <- NA
    smooth_cov_y <- NA
    sp <- NA
    time_cov_pred_grid <- NA
    time_cov_estim <- NA
    warning("covariance estimation did not succeed")
  } 
  
  gam1 <- NULL
  
  ###############
  # return output
  ###############
  out <- list(sigmasq = sigmasq, sigmasq_int = sigmasq_int, grid_mat_B = grid_mat_B, 
              grid_mat_C = grid_mat_C, grid_mat_E = grid_mat_E, 
              sp = sp, time_cov_estim = time_cov_estim, 
              time_cov_pred_grid = time_cov_pred_grid)
  return(out)
}

####################################################################################