##################################################################################################################
# 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
        }
      }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)
}

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