#!/bin/bash

usage="$(basename "$0") [-h] -r <reference> -i <fastq>

Calculate accuracy statistics for an assembly. 

    -h  show this help text.
    -r  reference, should be a fasta file. If correspondng bwa indices
        do not exist they will be created. (required).
    -i  fastq/a input assembly (required).
    -c  chunk size. Input reads/contigs will be broken into chunks
        prior to alignment, 0 will not chunk (default 100000).
    -C  catalogue errors. 
    -t  alignment threads (default: 1).
    -p  output file prefix (default: assm).
    -T  trim consensus to primary alignments of truth to assembly.
    -b  .bed file of reference regions to assess."

PREFIX="assm"
THREADS=1
CHUNK="100000"
rflag=false
iflag=false
bed_flag=false
catalogue_flag=false
trim_flag=false
ALIGN_OPTS=""
BED=""


while getopts ':hr:i:p:c:CTt:b:' option; do
  case "$option" in
    h  ) echo "$usage" >&2; exit;;
    r  ) rflag=true; REFERENCE=$OPTARG;;
    i  ) iflag=true; INPUT=$OPTARG;;
    p  ) PREFIX=$OPTARG;;
    c  ) CHUNK=$OPTARG;;
    C  ) catalogue_flag=true; ALIGN_OPTS="-m";;
    T  ) trim_flag=true;;
    t  ) THREADS=$OPTARG;;
    b  ) bed_flag=true; BED="--bed $OPTARG"; ALIGN_OPTS="-m";;
    \? ) echo "Invalid option: -${OPTARG}." >&2; exit 1;;
    :  ) echo "Option -$OPTARG requires an argument." >&2; exit 1;;
  esac
done
shift $(($OPTIND - 1))

if ! $iflag || ! $rflag; then
  echo "$usage" >&2;
  echo "-i and -r must be specified." >&2;
  exit 1;
fi

if (($CHUNK > 0)); then
    ALIGN_OPTS="$ALIGN_OPTS -c $CHUNK"
fi

if $trim_flag; then
    mini_align -i $REFERENCE -r $INPUT -p ${PREFIX}_trim -t $THREADS $ALIGN_OPTS
    ref_seqs_from_bam ${PREFIX}_trim.bam > ${PREFIX}_trim.fasta
    INPUT=${PREFIX}_trim.fasta
fi

mini_align -i $INPUT -r $REFERENCE -p $PREFIX -t $THREADS $ALIGN_OPTS 

STATS_THREADS=1
if [[ "$BED" != "" ]]; then
    STATS_THREADS=$THREADS
fi
stats_from_bam ${PREFIX}.bam -o ${PREFIX}_stats.txt -t ${STATS_THREADS} ${BED} 

summary_from_stats -i ${PREFIX}_stats.txt -pr -o ${PREFIX}_summ.txt

if $catalogue_flag; then
    OUTDIR=${PREFIX}_error_catalogue
    echo "Running catalogue_errors, saving data to ${OUTDIR}"
    catalogue_errors ${PREFIX}.bam -t $THREADS -o ${OUTDIR} ${BED}
fi

grep 'Percentage Errors' -A 7 -m 1 ${PREFIX}_summ.txt
grep 'Q Scores' -A 7 -m 1 ${PREFIX}_summ.txt

echo "All done, output written to ${PREFIX}_stats.txt and ${PREFIX}_summ.txt"
