module load anaconda
source activate snakemake

# Read input arguments. 
ref=$1 
bam=$2
ploidy=$3 
vcf=$4

# name gVCF file
gvcf=${vcf/.vcf.gz/.g.vcf.gz}

# if no GATK dictionary, index reference genome
if [ ! -f ${ref%.*}".dict" ] ; then
echo "GATK dictionary for $ref" >&2
picard CreateSequenceDictionary \
	REFERENCE=${ref} \
	OUTPUT=${ref%.*}".dict" 
fi

# If BAM index does not exist, index BAM.
if [ ! -s ${bam}.bai ]; then 
  echo 'indexing bam' 
  samtools index ${bam}
fi

# Call variants with GATK 4.1, output GVCF
gatk --java-options "-Xmx100g" HaplotypeCaller \
-R ${ref} \
-ploidy ${ploidy} \
-I ${bam} \
-ERC GVCF \
-O ${gvcf}

# Index gvcf 
gatk IndexFeatureFile \
  -I ${gvcf}

# GVCF to VCF. 
gatk --java-options '-Xmx100g' GenotypeGVCFs \
-R ${ref} \
--variant ${gvcf} \
-ploidy ${ploidy} \
--include-non-variant-sites true \
--output ${vcf}
# min base quality score is 10 by default.

#Error handling
if [ "$?" != "0" ]; then
	echo "[Error]" $LINENO "GATK failed!" 1>&2
	exit 1
fi

#### PRINT OUTPUT ####
echo "Done====" >&2

module load anaconda
source activate snakemake

# Read from command line: ref genome, fastq 1, fastq 2.
ref=${1} # 1st input is full path to reference genome
mapper=$2 # 2nd input is mapping algorithm 
p1=$3  # 3th input is full path to read 1
p2=$4  # 4th input is full path to read 2 (if paired-end)
bam=$5 # bam file.

# Define output directory. 
BAMS_DIR=$(dirname ${bam})

# Define temp directory.
TMP_DIR=${BAMS_DIR}/tmp/
mkdir -p ${TMP_DIR}

# Define prefix.
prefix=$(basename ${p1%%.*})

# Set names of intermediate files.
ref_index=${ref%.*}
ref_prefix=$(basename $ref_index)

# remove suffix
sam=${TMP_DIR}${prefix}_${mapper}_${ref_prefix}'.sam' 
rawbam=${TMP_DIR}${prefix}_${mapper}_${ref_prefix}'.bam' 
rgbam=${TMP_DIR}${prefix}_${mapper}_${ref_prefix}'_rg.bam' 
sortbam=${TMP_DIR}${prefix}_${mapper}_${ref_prefix}'_rg_sorted.bam'
echo $sam

# help message if no inputs are entered (-z checks if 1st arugment string is null).
if [ -z ${1} ]
then
  echo "This script will use Bowtie 2 to align two paired-end fastq files"
  echo "$(basename $0) <refGenome> <mapper> <dataset> <readFile1> <readFile2>"
  echo "First input is the Mtb ref genome"
  echo "Second input is the read mapper"
  echo "Third input is the dataset"
  echo "Forth input is location of file containing read 1"
  echo "Fifth input is location of file containing read 2 (if paired-end)" 
  exit
fi

# if no p2 given 
if [ -z ${p2} ]
then
  echo "read pair not specified"
fi

# get machine id_lane from SRA read identifier (old Illumina fastq format)
# if gzipped fastq
if [[ ${p1} == *.gz ]]; then
  seqid=$(zcat ${p1} | head -n 1)
else
# if not gzipped
  seqid=$(cat ${p1} | head -n 1)
fi
seqid="$(echo $seqid | cut -d' ' -f1)"
seqid="$(echo $seqid | cut -d':' -f3)"

id_lane=${seqid:-readgroup1} # default is "readgroup1"

#### MAPPING ####
# bowtie2 mapping
if [ $mapper == 'bowtie' ] || [ $mapper == 'bowtie2' ] ; then

  # if no indexing, index reference genome
  if [ ! -f ${ref%.*}".1.bt2" ] ; then
  echo "bowtie2 indexing $ref" >&2
  ${BOWTIE2_BUILD} ${ref} ${ref_index}
  fi 

  # map
  #if paired-end reads
  echo "mapping with bowtie2" >&2
  if [ ! -z ${p2} ]; then 
  bowtie2 --threads 7 -X 1100 -x ${ref_index} -1 ${p1} -2 ${p2} -S ${sam}
  # -x basename of reference index 
  # --un-gz gzips sam output
  # p is number of threads
  # end-to-end mapping is the default mode
  # -X 2000 increase maximum fragment length from default (500) to allow longer fragments (paired-end only) -X 2000 **used for roetzer data ***
  # ART simulations X = 1100 (mean fragment length = 650bp + 3 x 150-bp stdev)

  # if single-end reads
  elif [ -z ${p2} ]; then
    echo "single reads"
  bowtie2 --threads 7 -X 1100 -x ${ref_index} -U ${p1} -S ${sam}
  # -U for unpaired reads
  fi
  # Error handling
  if [ "$?" != "0" ]; then
    echo "[Error]" $LINENO "bowtie mapping failed ${p1}!" 1>&2
    exit 1
  fi
fi

# bwa mapping
if [ $mapper == 'bwa' ]; then
  # if no indexing, index reference genome
  if [ ! -f ${ref}".sa" ] ; then
  echo "bwa indexing $ref" >&2
  bwa index ${ref} 
  fi

  # map
  echo "mapping with bwa" >&2
  # if paired-end reads
  if [ ! -z ${p2} ]; then 
    bwa mem -t 7 ${ref} ${p1} ${p2} > ${sam}
   # -t no. of threads.
  # if single-end reads
  elif [ -z ${p2} ]; then
      echo "single reads"
    bwa mem -t 7 ${ref} ${p1} >  ${sam}
  fi
  # Error handling
  if [ "$?" != "0" ]; then
    echo "[Error]" $LINENO "bwa mapping failed ${p1}!" 1>&2
    exit 1
  fi
fi

### POST-PROCESSING ####

# Convert sam to bam 
sambamba view -t 7 -S -h ${sam} -f bam -o ${rawbam}
# -S auto-detects input format, -h includes header, -o directs output

# Add/replace read groups for post-processing with GATK
picard AddOrReplaceReadGroups \
  INPUT=${rawbam} \
  OUTPUT=${rgbam} \
  RGID=${id_lane} \
  RGLB=library1 \
  RGPU=${id_lane} \
  RGPL="illumina" \
  RGSM=${prefix}

# Sort the BAM 
sambamba sort ${rgbam}

# Index BAM
sambamba sort ${rgbam} index -t7 --out ${sortbam}

# Remove duplicates. (-r = remove)
sambamba markdup -r -p -t7 ${sortbam} ${bam} --tmpdir=${TMP_DIR} 

# Remove intermediate files.
#rm ${sam} ${rawbam} ${rgbam} #${sortbam} ${sortbam}.bai

# Error handling
if [ "$?" != "0" ]; then
    echo "[Error]" $LINENO "Remove dups failed ${p1}!" 1>&2
    exit 1
fi

### PRINT OUTPUT ####
echo "Done====" >&2

module purge
module load anaconda
source activate snakemake

# Kraken database.
DBNAME=/labs/jandr/walter/varcal/data/refs/kraken/

# Read from command line paired end 1 and paired end 2 - trimmed, zipped fastq files.
p1=$1
p2=$2
filt1=$3
filt2=$4
report=$5

# name outputs
prefix=$(basename ${p1/_trim_1.*})
LOGDIR=$(dirname $report)/

# run kraken to taxonomically classify paired-end reads and write output file.
kraken2 --db ${DBNAME} --paired --gzip-compressed --threads 8 --report ${report} --use-names ${p1} ${p2} --output ${LOGDIR}${prefix}.out

# select for reads directly assigned to Mycobacterium genus (G) (taxid 1763), reads assigned directly to Mycobacterium tuberculosis complex (G1) (taxid 77643), and reads assigned to Mycobacterium tuberculosis (S) and children. *this includes reads assigned to Mtb and those assigned to the genus, but not a different species.

grep -E 'Mycobacterium \(taxid 1763\)|Mycobacterium tuberculosis' ${LOGDIR}${prefix}.out | awk '{print $2}' > ${LOGDIR}${prefix}_reads.list

# use bbmap to select reads corresponding to taxa of interest.
filterbyname.sh  int=false in1=${p1} in2=${p2} out1=${filt1} out2=${filt2} names=${LOGDIR}${prefix}_reads.list include=true overwrite=true

# test if kraken filtered fastq files are the same length
len1=$(cat ${filt1}  | wc -l)
len2=$(cat ${filt2} | wc -l)  

echo $len1 $len2

# Write error to log.
if [ $len1 != $len2 ]; then
  echo ${prefix} error: paired-end files are of different lengths. 
fi  

# Summarize Kraken statistics. 
/labs/jandr/walter/tb/scripts/kraken_stats.sh ${report} 

# Error handling
if [ "$?" != "0" ]; then
	echo "[Error]" $LINENO "Kraken ${p1} failed!" 1>&2
	exit 1
fi

module load anaconda
source activate snakemake

# Read from command line: ref genome, fastq 1, fastq 2.
p1=$1  # fastq 1
p2=$2  # fast1 2
trim1=$3 # trimmed/adapter-removed fastq 1
trim2=$4 # trimmed/adapter-removed fastq 2

# Print input and output filenames.
echo $p1 $p2 
echo 'output files' $trim1 $trim2

# Make temp directory for intermediate files. 
mkdir -p $(dirname $p1)/tmp
TMP_DIR=$(dirname $p1)/tmp/

# Name temporary files.
tmp1=${TMP_DIR}$(basename ${p1/.fastq.gz/_val_1.fq.gz})
tmp2=${TMP_DIR}$(basename ${p2/.fastq.gz/_val_2.fq.gz})


#### Trim READS ####
echo "trimming reads"

# trip adapters 
trim_galore  --gzip --stringency 3 --paired ${p1} ${p2} --output_dir ${TMP_DIR} 2>&1 
# -q Trims low-quality (Phred scaled base quality < 20) in addition to adapter removal
# --length Sets min read length at 70; bwa mem is optimized for reads >=70 bp. (Abbott)

# second, try cutadapt to remove next seq poly-G's
cutadapt -f 'fastq' --nextseq-trim=20  --minimum-length=20 --pair-filter=any -o ${trim1} -p ${trim2} ${tmp1} ${tmp2} 2>&1 
# removes poly-G tails -NovaSeq. In those instruments, a ‘dark cycle’ (with no detected color) encodes a G. However, dark cycles also occur when when sequencing “falls off” the end of the fragment. The read then contains a run of high-quality, but incorrect `G` calls 

# run fastqc again
#/srv/gsfs0/software/fastqc/0.11.4/fastqc --noextract --threads 8 --outdir ${QC_DIR} ${cutadapt1} 
#/srv/gsfs0/software/fastqc/0.11.4/fastqc --noextract --threads 8 --outdir ${QC_DIR}  ${cutadapt2}

# Test if both paired-end reads have same number of lines.
len1=$(zcat $trim1 | wc -l)
len2=$(zcat $trim2 | wc -l)

# Create error if two files are different lengths.
if [ $len1 != $len2 ]; then   
   echo ${prefix}'_1.fq'  
   error: paired-end files are of different lengths. 
fi  

# remove extra files
rm ${tmp1} ${tmp2}

# Error handling
if [ "$?" != "0" ]; then
	echo "[Error]" $LINENO "Read trimming ${p1} failed!" 1>&2
	exit 1
fi

module load bcftools

# read from command line.
ref=$1
vcf=$2
sample_name=$3
bed=$4
fasta=$5

# Get sample name for correct genotype
samp=$(bcftools query -l ${vcf} )

# Make consensus masked sequence & rename fasta header. 
bcftools consensus --include 'TYPE!="indel"' --mask ${bed} --fasta-ref ${ref} --sample ${samp} --absent 'N' --missing 'N' ${vcf}  | \
    seqtk rename - ${sample_name}  > ${fasta}

shell:    
  '''
  # Coverage summary statistics
  workflow/scripts/cov_stats.sh {input.ref_path} {input.combined_bam} {output.cov_stats} &>> {log}
  # Mosdepth coverage along genome (for plotting)
  mosdepth --by 2000 -Q 30 {params.prefix} {input.combined_bam}
  '''

shell:
  '''
  # Rename Chromosome to be consistent with snpEff/Ensembl genomes.
  zcat {input.vcf}| sed 's/NC_000962.3/Chromosome/g' | bgzip > {output.rename_vcf}
  tabix {output.rename_vcf}

  # Run snpEff
  java -jar -Xmx8g {config[snpeff]} eff {config[snpeff_db]} {output.rename_vcf} -dataDir {config[snpeff_datapath]} -noStats -no-downstream -no-upstream -canon > {output.tmp_vcf}

  # Also use bed file to annotate vcf
  bcftools annotate -a {params.bed} -h {params.vcf_header} -c CHROM,FROM,TO,FORMAT/PPE {output.tmp_vcf} > {output.ann_vcf}

  '''