npoint

#!/usr/bin/env Rscript

#### Make sure packages are installed and load them 

#rniftilib is not in this location
#if (!suppressWarnings(suppressPackageStartupMessages(require(Rniftilib)))) install.packages("Rniftilib", repos="http://R-Forge.R-project.org")
if (!suppressWarnings(suppressPackageStartupMessages(require(argparse)))) install.packages("argparse")
if (!suppressWarnings(suppressPackageStartupMessages(require(doParallel)))) install.packages("doParallel")
suppressWarnings(suppressPackageStartupMessages(library(RNifti)))
suppressWarnings(suppressPackageStartupMessages(library(argparse)))
suppressWarnings(suppressPackageStartupMessages(library(doParallel)))
suppressWarnings(suppressPackageStartupMessages(library(neuropointillist)))

#### Take in command line arguments 

parser <- ArgumentParser(description="This program prepares your MRI data for group-level mixed effects modeling")

parser$add_argument("-m", "--mask", nargs=1, type="character", help="Mask limiting the voxels that will be analyzed", required=TRUE)
parser$add_argument("--set1", nargs=1,  help="List of files at first occasion", required=TRUE)
parser$add_argument("--set2", nargs=1, help="List of files at second occasion")
parser$add_argument("--set3", nargs=1, help="List of files at third occasion")
parser$add_argument("--set4", nargs=1, help="List of files at fourth occasion")
parser$add_argument("--set5", nargs=1, help="List of files at fifth occasion")
parser$add_argument("--setlabels1", nargs=1, help="Covariates for files at first occasion", required=TRUE)
parser$add_argument("--setlabels2", nargs=1, help="Covariates for files at second occasion")
parser$add_argument("--setlabels3", nargs=1, help="Covariates for files at third occasion")
parser$add_argument("--setlabels4", nargs=1,help="Covariates for files at fourth occasion")
parser$add_argument("--setlabels5", nargs=1, help="Covariates for files at fifth occasion")
parser$add_argument("--model", nargs=1, help="R code that defines the voxelwise-model and any initialization", required=TRUE)
parser$add_argument("--covariates", nargs=1, type="character", help="Covariates that will be merged with the design matrix")
parser$add_argument("--output", nargs=1, type="character", help="Output prefix to prepend to output files", required=TRUE)
parser$add_argument("--debugfile", nargs=1, type="character", help="Save voxeldat and designmat objects to this file to develop, test and debug the processVoxel function")
parser$add_argument("-t", "--testvoxel", type="integer", help="Specify a voxel on which the model works to determine output files", default="-1")
parser$add_argument("-p", "--processors", type="integer", help="Run using shared memory with p processors")
parser$add_argument("--sgeN",  type="integer", nargs=1, help="Run using SGE generating N jobs")
parser$add_argument("--slurmN",  type="integer", nargs=1, help="Run using Slurm generating N jobs")
parser$add_argument("--pbsN",  type="integer", nargs=1, help="Run using PBS generating N jobs")
parser$add_argument("--pbsPre",  type="character", nargs=1, help="Name of PBS preamble file to use if you wish to override default settings")
parser$add_argument("--permute",  type="integer", nargs=1, help="Enter permutation testing mode with N permutations. This will generate N jobs, ignoring the number of jobs passed to any scheduler parameters. Each permutation will create a single 3D output file.")

if (file.exists("readargs.R")) {
    source("readargs.R")
    if (exists("cmdargs")) { # check that cmdargs is defined
        args <- parser$parse_args(cmdargs)
    } else {
        args <- parser$parse_args()
    }
} else {
    args <- parser$parse_args()
}
    

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

#### Check for mask and read it. It is mandatory and must exist.
maskfile <- args$mask
tryCatch({
    mask <- readNifti(maskfile);
}, error=function(e) {
    cat("Could not read mask file: ", maskfile, "\n")
    stop(e)
})

# save mask dimensions
mask.dims <- dim(mask)

# reduce to vector and obtain list of nonzero vertices and the x,y,x
mask.vector <- as.vector(mask)
mask.vertices <- which(mask.vector > 0)

# assemble the indices to create a reverse lookup table
mask.arrayindices <- data.frame(which(mask > 0, arr.in=TRUE))
mask.arrayindices$vertex <- mask.vertices

#### Save original arguments for writing out calling info
origargs <- args


#### Do argument checking
args <- npointCheckArguments(args)

#### Are we running in parallel?
if (!is.null(args$processors) || !is.null(args$sgeN) || !is.null(args$slurmN) || !is.null(args$pbsN)) {
    runningParallel =TRUE
} else {
    runningParallel= FALSE
}

#### Did someone specify some combination of sge, slurm and PBS? Exit if they did
multipleschedulerflags <- sum(!is.null(args$sgeN), !is.null(args$slurmN), !is.null(args$pbsN))

if (multipleschedulerflags > 1) {
    message("You tried to create driver files for multiple schedulers. You can only specify one at a time.")
    message("If you want Slurm, use only the --slurmN flag.")
    message("If you want PBS, use only the --pbsN flag.")
    message("If you want SGE, use only the --sgeN flag.")        
    stop("Please try again specifying only one scheduler flag.")
}

#### Check to see that PBS preamble is only specified with PBS option
if (!is.null(args$pbsPre)) {
    if (is.null(args$pbsN)) {
        stop("You have specified a PBS preamble without the PBS scheduler flag. You cannot use a PBS preamble with other schedulers.")
    }
    # make sure the preamble file exists
    if (!file.exists(args$pbsPre)) {
        stop("PBS preamble file ", args$pbsPre, " does not exist!")
    }
}

#### If we are writing out files for a scheduler, identify the number of files
nSchedulerJobs <- -1 # for sanity
if (runningParallel) {
    if (!is.null(args$sgeN)) {
        nSchedulerJobs <- args$sgeN
    } else if (!is.null(args$slurmN)) {
        nSchedulerJobs <- args$slurmN
    } else if (!is.null(args$pbsN)) {
        nSchedulerJobs <- args$pbsN
    }

}

#### If we are in permutation mode, reset the number of scheduler jobs 
#### In this mode, we will not split the data at all. We will generate a makefile
#### that runs the processVoxel code across all the data to generate each permutation
permutationMode <- 0
if (!is.null(args$permute)) {
    permutationMode <- 1
    nSchedulerJobs <- args$permute
    if (!is.null(args$processors)) {
        stop("Cannot specify -p flag with permute mode. You must specify a scheduler with some number of jobs, which will be ignored, and one job will be created for each permuation.")
    }
    if (!runningParallel) {
        stop("If running in permute mode, you must specify a scheduler with some number of jobs. This number of jobs will be ignored, and one job will be created for each permuation.")
        }
}

#### A lookup function to convert FSL indices to a vertex number for testing
imagecoordtovertex <- function(x,y,z) {
    # first add 1 to convert from index at zero to index at 1
    nx <- x+1
    ny <- y+1
    nz <- z+1
    row <- mask.arrayindices[which(mask.arrayindices$dim1==nx&
                                   mask.arrayindices$dim2==ny&
                                   mask.arrayindices$dim3==nz),]
    if(is.data.frame(row) && nrow(row)==0) {
        warning("This coordinate is not in the mask; returning 1")
        return(1)
    } else {
        return(as.integer(row.names(row)))
    }
}
    

    
###############################################################################
#### Calculate the number of data sets
numberdatasets <- sum(!is.null(args$set1),
                          !is.null(args$set2),
                          !is.null(args$set3),
                          !is.null(args$set4),
                          !is.null(args$set5))

###############################################################################
#### Read in all the data sets

cat("Reading", numberdatasets, "data sets.\n")
data <- npointReadDataSets(args,numberdatasets,mask.vertices);
voxeldat <- data$voxeldat
designmat <-data$designmat
rm(data)
gc()

###############################################################################
### Create the output directory if it does not exist

dir <- dirname(args$output)
if (!dir.exists(dir)) {
    dir.create(dir, recursive=TRUE)
}

###############################################################################
#### Read in covariates if specified and merge with other covariates specified
#### on the command line
if (!is.null(args$covariates)) {
    designmat <- npointMergeDesignmatWithCovariates(designmat,args$covariates,dim(voxeldat)[1])
}

###############################################################################
### If debugging file is specified, save out design matrix and voxel matrix
### to this file
if(!is.null(args$debugfile)) {
    dir <- dirname(args$output)
    # add the prefix to the debug file name
    debugfilename <- paste(dir, args$debugfile,sep="/")
    save(designmat,voxeldat, imagecoordtovertex, mask.arrayindices, file=debugfilename)
}

###############################################################################
#### read model code
if (!is.null(args$model)) {
    modelfile <- args$model
    if (!file.exists(modelfile)) {
        stop("model file ", modelfile, " does not exist!")
    }
    result <- tryCatch({
        source(modelfile)
    }, error=function(e) {
        cat("There were errors in the model file ", modelfile, "\n")
        stop(e)

    })
}
                   
# check to see that processVoxel is defined
if(!exists('processVoxel')) {
    stop("The model file did not define the processVoxel function")
}


###############################################################################
#### Do the parallel processing

nvertices <- length(mask.vertices)
if (permutationMode) { # we are using SGE or SLURM or PBS
    # save the design matrix
    designmatname <- paste(args$output, "designmat.rds", sep="")
    makefilename <- paste(dirname(args$output), "/Makefile", sep="")
    nextflowfilename <- paste(dirname(args$output), "/make.nf", sep="")
    masterscript.local <- paste(dirname(args$output), "/runme.local", sep="")
    saveRDS(designmat,designmatname)
    attach(designmat) # attach to the designmat
                                        # test one voxel to obtain names for return vals
    if (args$testvoxel < 0) {
        args$testvoxel <- trunc(dim(voxeldat)[2]/2)
    }
    permutationNumber <- 1 # set permutation number
    tryCatch({ out <- processVoxel(args$testvoxel)},error=function(e) {
        message("error testing the model on a random voxel to determine output filenames")
        message("try providing the -t option to give a specific voxel for testing")
        message(e)
        stop("Exiting before generating makefile.")})
    npointWriteMakefilePermute(basename(args$output), names(out), paste(getwd(), "/",args$model,sep=""), basename(designmatname), makefilename, masterscript.local, nSchedulerJobs)
    npointWriteNextflowPermute(basename(args$output), names(out), paste(getwd(), "/",args$model,sep=""), basename(designmatname), nextflowfilename, nSchedulerJobs)
    npointSplitDataSize(dim(voxeldat)[2],voxeldat,args$output,mask)

    if (!is.null(args$sgeN)) {
        masterscript.scheduler <- paste(dirname(args$output), "/runme.sge", sep="")
        jobscript <- paste(dirname(args$output), "/sgejob.bash", sep="")
 #       npointWriteSGEsubmitscript(basename(args$output), names(out), paste(getwd(), "/",args$model,sep=""), basename(designmatname), masterscript.scheduler,jobscript, njobs)
        
    } else if (!is.null(args$slurmN)) {
        masterscript.scheduler <- paste(dirname(args$output), "/runme.slurm", sep="")
        jobscript <- paste(dirname(args$output), "/slurmjob.bash", sep="")
  #      npointWriteSlurmsubmitscript(basename(args$output), names(out), paste(getwd(), "/",args$model,sep=""), basename(designmatname), masterscript.scheduler,jobscript, njobs)              
        
    } else if (!is.null(args$pbsN)) {
        masterscript.scheduler <- paste(dirname(args$output), "/runme.pbs", sep="")
        jobscript <- paste(dirname(args$output), "/pbsjob.bash", sep="")
   #     npointWritePBSsubmitscript(basename(args$output), names(out), paste(getwd(), "/",args$model,sep=""), basename(designmatname), masterscript.scheduler,jobscript, njobs, args$pbsPre)              
    } else {              
         # We better not get here
        stop("Cannot identify which scheduler to output")
    }
    npointWriteCallingInfo(origargs)

} else {
    if (runningParallel) {
        if (!is.null(args$processors)) {
            attach(designmat) # we attach to the designmat
            cl <- makeCluster(args$processors, type="FORK")
            cat("Exporting data to cluster.\n")
            clusterExport(cl,varlist=c("voxeldat"))
            cat("Starting parallel job using", args$processors, "cores.\n")
            cat("Use top to make sure that no threads are using more than 100% of the CPU.\n")
            system.time(results <-parSapply(cl,1:nvertices, processVoxel))
            stopCluster(cl)
            npointWriteOutputFiles(args$output,results,mask)
            npointWriteCallingInfo(origargs)
        } else { # we are using SGE or SLURM or PBS
                                        # save the design matrix
            designmatname <- paste(args$output, "designmat.rds", sep="")
            makefilename <- paste(dirname(args$output), "/Makefile", sep="")
            masterscript.local <- paste(dirname(args$output), "/runme.local", sep="")
            saveRDS(designmat,designmatname)
            attach(designmat) # attach to the designmat
                                        # test one voxel to obtain names for return vals
            if (args$testvoxel < 0) {
                args$testvoxel <- trunc(dim(voxeldat)[2]/2)
            }
            tryCatch({ out <- processVoxel(args$testvoxel)},error=function(e) {
                message("error testing the model on a random voxel to determine output filenames")
                message("try providing the -t option to give a specific voxel for testing")
                message(e)
                stop("Exiting before generating makefile.")})
            npointWriteMakefile(basename(args$output), names(out), paste(getwd(), "/",args$model,sep=""), basename(designmatname), makefilename, masterscript.local)


       # do sanity check to make sure we obtained the number of jobs
        if (nSchedulerJobs < 0) {
            stop("Number of scheduler jobs incorrectly specified.")
        }
        # split up the data into chunks and write out scripts to process
        if (nSchedulerJobs > nvertices) {
            stop("Number of scheduler jobs requested is greater than the number of vertices")
        } else {
          cat("no. of vertices", nvertices, "\n")
          #sgeN * size is now larger than the # of vertices
  	  size <- ceiling(nvertices/nSchedulerJobs)
          cat("size", trunc(nvertices/nSchedulerJobs), "\n")
          njobs <- npointSplitDataSize(size,voxeldat,args$output,mask)
          cat("no. of jobs", njobs ,"\n")


          if (!is.null(args$sgeN)) {
              masterscript.scheduler <- paste(dirname(args$output), "/runme.sge", sep="")
              jobscript <- paste(dirname(args$output), "/sgejob.bash", sep="")
              npointWriteSGEsubmitscript(basename(args$output), names(out), paste(getwd(), "/",args$model,sep=""), basename(designmatname), masterscript.scheduler,jobscript, njobs)
            
          } else if (!is.null(args$slurmN)) {
              masterscript.scheduler <- paste(dirname(args$output), "/runme.slurm", sep="")
              jobscript <- paste(dirname(args$output), "/slurmjob.bash", sep="")
              npointWriteSlurmsubmitscript(basename(args$output), names(out), paste(getwd(), "/",args$model,sep=""), basename(designmatname), masterscript.scheduler,jobscript, njobs)              

          } else if (!is.null(args$pbsN)) {
              masterscript.scheduler <- paste(dirname(args$output), "/runme.pbs", sep="")
              jobscript <- paste(dirname(args$output), "/pbsjob.bash", sep="")
              npointWritePBSsubmitscript(basename(args$output), names(out), paste(getwd(), "/",args$model,sep=""), basename(designmatname), masterscript.scheduler,jobscript, njobs, args$pbsPre)              
          } else {              
              # We better not get here
              stop("Cannot identify which scheduler to output")
          }
          npointWriteCallingInfo(origargs)
      }
    }
    } else {
        
    cat("Starting sequential job\n")
    cat("You might want to check whether your model is multithreaded\n")
    cat("because your code might run faster if you limit the number of threads\n")
    system.time(results <-sapply(1:nvertices, processVoxel))
    npointWriteOutputFiles(args$output,results,mask)
    npointWriteCallingInfo(origargs)
}
}