#!/bin/bash

#SBATCH --nodes=1
#SBATCH --cpus-per-task=8
#SBATCH --mem-per-cpu=4GB
#SBATCH --time=8:00:00
#SBATCH --qos=avram-b
#SBATCH --output=DESEq2.%j.out
#SBATCH --job-name=DESeq2

##### LAST UPDATED 2017.11.21 KJL #####

##### ARGUMENTS AND USAGE #####
usage="$(basename $0) [-h] [-d dir] [-r file] [-g file] [-t string] [-c string] bam_list -- Program to run DESeq2 differential analysis on RNA-Seq bam files.

Where: bam_list is a space-separated list of RNA-Seq bam files (e.g. RNA_WT_Th0_sample1.bam RNA_WT_Th0_sample2.bam RNA_WT_Th2_sample1.bam RNA_WT_Th2_sample2.bam)

Optional Arguments:
    -h  Show this help text
    -d  Set target directory (directory to write analysis files) (default: $PWD/analysis/DESeq2)
    -r  Set RNA-Seq sample list file (default: \$HOME/lib/RNA_samples.txt)
    -g  Set annotated GTF file (default: \$HOME/lib/mm10_with_gene_ids.gtf)
    -t  Set test condition (column from RNA-Seq sample list file) (i.e. Genotype or Cell_subset) (default: Genotype)
    -c  Set control group (i.e. WT if test condition is Genotype, or Th0 if test condition is Cell_subset (default: WT)

Example usage:
    sbatch --chdir=/ufrc/avram/kyle.lorentsen/RNA/CD8_RNA_Seq_Run --export=t=Cell_subset,c=Th0,bam_list=\"RNA_WT_Th0_sample1.bam RNA_WT_Th0_sample2.bam RNA_WT_Th2_sample1.bam RNA_WT_Th2_sample2.bam\" $HOME/bin/$(basename $0)

    or

    $(basename $0) -d ufrc/avram/kyle.lorentsen/RNA/CD8_RNA_Seq_Run -c Th0 -t Cell_subset RNA_WT_Th0_sample1.bam RNA_WT_Th0_sample2.bam RNA_WT_Th2_sample1.bam RNA_WT_Th2_sample2.bam"

while getopts ':hd:r:g:t:c:' option; do
    case "$option" in
    h)  echo "$usage"
        exit 0
        ;;
	d)  if [[ -d "$OPTARG" ]]; then
            d="$OPTARG"
            printf "Target directory manually set to %s\n" "$OPTARG"
        else
            printf "%s doesn't exist or is not a valid directory. Please specify a valid directory and rerun %s. Exiting...\n" "$OPTARG" "$(basename $0)"
            exit 1
        fi
        ;;
    r)  if [[ -s "$OPTARG" ]]; then
            r="$OPTARG"
            printf "RNA_Seq sample list file manually set to: %s\n" "$OPTARG"
        else
            printf "%s is empty or does not exist. Please specify valid RNA_Seq sample list file and rerun %s. Exiting...\n" "$OPTARG" "$(basename $0)"
        fi
        ;;
    g)  g="$OPTARG"
        printf "Annotated GTF file manually set to %s\n" "$OPTARG"
        ;;
    t)  t="$OPTARG"
        printf "Test condition manually set to: %s\n" "$OPTARG"
        ;;
    c)  c="$OPTARG"
        printf "Control group manually set to: %s\n" "$OPTARG"
        ;;
   \?)  printf "Invalid option: -%s.\n" "$OPTARG"
        echo "$usage"
        exit 1
        ;;
    :)  printf "Invalid option: -%s requires an argument.\n" "$OPTARG"
        echo "$usage"
        exit 1
        ;;
    esac
done
shift $((OPTIND -1))

date;pwd

##### LOAD MODULES #####
module load R

##### SET VARIABLES #####
target_dir="${d:-${PWD}/analysis/DESeq2}"
if [[ "$target_dir" == "$PWD/analysis/DESeq2" && ! -d "$target_dir" ]]; then # IF TARGET DIR DOES NOT EXIST
    mkdir -p "$target_dir"
elif [[ ! -d "$target_dir" ]]; then # THIS CAN ONLY BE TRUE IF SOMEONE CALLED THIS FUNCTION THROUGH SBATCH AND DIRECTORY EXPORTED AN INVALID "d" variable. (e.g. sbatch --export=d=/some/invalid/dir DESeq2.sh)
    printf "%s doesn't exist or is not a valid directory. Please specify a valid directory and rerun %s. Exiting...\n" "$target_dir" "$(basename $0)"
    exit 1
fi

gtf="${g-${HOME}/lib/mm10_with_gene_ids.gtf}"
[[ -f "$gtf" ]] || { printf "%s does not exist.  Please specify valid annotated GTF file and rerun %s. Exiting...\n" "$gtf" "$(basename $0)"; exit 1; } # VERIFY ANNOTATED GTF FILE IS VALID

RNA_samples="${r:-${HOME}/lib/RNA_samples.txt}"
[[ -s "$RNA_samples" ]] || { printf "%s is empty or does not exist. Please specify valid RNA_Seq sample list file and rerun %s. Exiting...\n" "$OPTARG" "$(basename $0)"; exit 1; } # VERIFY RNA SAMPLE LIST FILE IS VALID

condition="${t:-Genotype}"
read -ra condition_list <<<$(head -n1 "$RNA_samples") # GENERATE ARRAY OF POSSIBLE CONDITIONS
[[ "${condition_list[@]}" =~ "$condition" ]] || { printf "%s is not a valid condition.  Please select a condition from the following list: %s. Exiting...\n" "$condition" "${condition_list[*]}"; exit 1; } #VERIFY SELECTED CONDITION IS VALID

baseline="${c:-WT}"
baseline_check=$(awk -F$'\t' -v condition=$condition -v baseline=$baseline 'BEGIN {s=0} 
    NR==1{for(i=1; i<=NF; i++) if($i==condition) {a[i]++; break} }
    {for (i in a) if ($i==baseline) {s++; exit} }
    END {print s}' "$RNA_samples")
[[ $baseline_check > 0 ]] || { printf "There are no RNA samples of type \"%s\" under condition \"%s\". Please select a valid control group and rerun %s. Exiting...\n" "$baseline" "$condition" "$(basename $0)"; exit 1; } #VERIFY THAT AT LEAST ONE SAMPLE UNDER SELECTED CONDITION MATCHES BASELINE TYPE

#ADD TEST TO CHECK IF AT LEAST ONE SAMPLE (BUT NOT ALL SAMPLES) HAVE BASELINE AS VALUE UNDER CONDITION
bam_list="${bam_list:-$@}"
if [[ -z "$bam_list" ]]; then
    printf "No list of BAM files detected. Please provide valid list of bam files and rerun %s. Exiting...\n" "$(basename $0)"
    exit 1
else
    IFS=' ' read -ra bam_array<<<"$bam_list"
    [[ "${#bam_array[@]}" > 1 ]] || { printf "At least two BAM files must be provided. Please update BAM file list and rerun %s. Exiting...\n" "$(basename $0)"; exit 1; } # VERIFY THAT AT LEAST TWO BAM FILES HAVE BEEN PROVIDED
    for bam in "${bam_array[@]}"; do # VERIFY THAT ALL THE GIVEN BAM FILES EXIST
        [[ -f "$bam" ]] || { printf "%s can not be found.  Please give valid bam file list and rerun %s. Exiting...\n" "$bam" "$(basename $0)"; exit 1; } 
    done

    uniqueBams=$(printf "%s\n" "${bam_array[@]}" | awk '!($0 in seen){seen[$0];c++} END {print c}')
    [[ $uniqueBams == "${#bam_array[@]}" ]] || { printf "Duplicate BAM files detected in given list.  Each BAM file can only be given once, please remove duplicates and rerun $0. Exiting...\n" "$(basename $0)"; exit 1; } # VERIFY THAT THERE ARE NO DUPLICATE BAM FILES IN GIVEN BAM LIST

fi

##### BEGIN SCRIPT #####
printf "Condition = %s\n" "$condition"
printf "Baseline = %s\n" "$baseline"
printf "RNA Sample list file = %s\n" "$RNA_samples"
printf "Annotated GTF file = %s\n" "$gtf"
printf "Bam_list = %s\n" "$bam_list"

( cd "$target_dir" && Rscript "$HOME"/bin/DESeq2.R -c "$condition" -b "$baseline" -r "$RNA_samples" -g "$gtf" -l "$bam_list" ) # RUN DESEQ2.R WITHIN SUBSHELL TO CHANGE DIRECTORY TO ANALYSIS DIR

( cd "$target_dir" && Rscript "$HOME"/bin/annotate_genes.R -i DESeq2_analysis.txt -o DESeq2_analysis_annotated.txt ) # RUN ANNOTATE_GENES_DESE2.R WITHIN SUBSHELL TO CHANGE DIRECTORY TO ANALYSIS DIR

date
