DAJIN

#!/bin/bash

################################################################################
# Initialize shell environment
################################################################################

set -u
umask 0022
export LC_ALL=C

mkdir -p .DAJIN_temp

################################################################################
# Define the functions for printing usage and error message
################################################################################

VERSION=0.7.0

usage() {
  cat <<-USAGE
Usage     : DAJIN -i [text file] (described at "Input")

Example   : ./DAJIN/DAJIN -i DAJIN/example/design.txt

Input     : Input file should be formatted as below:
            # Example
            ------
            design=DAJIN/example/example.fa
            input_dir=DAJIN/example/fastq
            control=barcode01
            genome=mm10
            grna=CCTGTCCAGAGTGGGAGATAGCC,CCACTGCTAGCTGTGGGTAACCC
            output_dir=DAJIN_example
            threads=10
            filter=on
            ------
            - design: a multi-FASTA file contains sequences of each genotype. ">wt" and ">target" must be included.
            - input_dir: a directory contains FASTA or FASTQ files of long-read sequencing
            - control: control barcode ID
            - grna: gRNA sequence(s) including PAM. multiple gRNA sequences must be deliminated by comma.
            - genome (optional): reference genome. e.g. mm10, hg38
            - output_dir (optional): output directory name. optional. Default is "DAJIN_results"
            - threads (optional; integer): Default is two-thirds of available CPU threads.
            - filter (optional; "on" or "off"): set filter to remove very minor allele (less than 3%). Default is "on"
USAGE
}

usage_and_exit() {
  usage
  exit 1
}

error_exit() {
  echo "$@" 1>&2
  exit 1
}

################################################################################
# Parse arguments
################################################################################
[ $# -eq 0 ] && usage_and_exit

if ! [ -r "$2" ]; then
  error_exit "$2: No such file"
fi

cat "$2" | tr -d "\r" >.DAJIN_temp/design.txt

while [ $# -gt 0 ]; do
  case "$1" in
  --help | --hel | --he | --h | '--?' | -help | -hel | -he | -h | '-?')
    usage_and_exit
    ;;
  --version | --versio | --versi | --vers | --ver | --ve | --v | \
    -version | -versio | -versi | -vers | -ver | -ve | -v)
    echo "DAJIN version: $VERSION" && exit 0
    ;;
  --input | --in | --i | -i)
    design=$(cat .DAJIN_temp/design.txt | grep "^design" | sed -e "s/ //g" -e "s/.*=//g")
    input_dir=$(cat .DAJIN_temp/design.txt | grep "^input_dir" | sed -e "s/ //g" -e "s/.*=//g")
    control=$(cat .DAJIN_temp/design.txt | grep "^control" | sed -e "s/ //g" -e "s/.*=//g")
    genome=$(cat .DAJIN_temp/design.txt | grep "^genome" | sed -e "s/ //g" -e "s/.*=//g")
    grna=$(cat .DAJIN_temp/design.txt | grep "^grna" | sed -e "s/ //g" -e "s/.*=//g")
    output_dir=$(cat .DAJIN_temp/design.txt | grep "^output_dir" | sed -e "s/ //g" -e "s/.*=//g")
    threads=$(cat .DAJIN_temp/design.txt | grep "^threads" | sed -e "s/ //g" -e "s/.*=//g")
    filter=$(cat .DAJIN_temp/design.txt | grep "^filter" | sed -e "s/ //g" -e "s/.*=//g")
    TEST=$(cat .DAJIN_temp/design.txt | grep "^TEST" | sed -e "s/ //g" -e "s/.*=//g")
    ;;
  -*)
    error_exit "Unrecognized option : $1"
    ;;
  *)
    break
    ;;
  esac
  shift
done

#===========================================================
# Check required arguments
#===========================================================

[ -z "$design" ] && error_exit "design argument is not specified"
[ -z "$input_dir" ] && error_exit "input_dir argument is not specified"
[ -z "$control" ] && error_exit "control argument is not specified"
[ -z "$grna" ] && error_exit "grna argument is not specified"

#===========================================================
# Check fasta file
#===========================================================

[ -f "$design" ] || error_exit "$design: No such file"

[ "$(grep -c -e '>wt' -e '>target' ${design})" -ne 2 ] &&
  error_exit "$design: design must include '>target' and '>wt'. "

#===========================================================
# Check directory
#===========================================================

[ -d "${input_dir}" ] || error_exit "$input_dir: No such directory"

fastq_num=$(find "${input_dir}"/* -type f | grep -c -e ".fq" -e ".fastq")

[ "$fastq_num" -eq 0 ] && error_exit "$input_dir: No FASTQ file in directory"

#===========================================================
# Check control
#===========================================================

if find "${input_dir}"/ -type f | grep -q "${control}"; then
  :
else
  error_exit "$control: No control file in ${input_dir}"
fi

#===========================================================
# Check genome
#===========================================================

genome_check=$(
  wget -q -O - "http://hgdownload.soe.ucsc.edu/downloads.html" |
    grep hgTracks |
    grep -c "${genome:-XXX}"
)

#===========================================================
# Check grna
#===========================================================

set $(echo "${grna}" | sed "s/,/ /g")
x=1
while [ "${x}" -le $# ]; do
  [ "$(grep -c ${1} ${design})" -eq 0 ] && error_exit "No gRNA sites"
  x=$(("${x}" + 1))
done

#===========================================================
# Check output directory name
#===========================================================

mkdir -p .DAJIN_temp/"$output_dir" 2>/dev/null ||
  error_exit "$output_dir: invalid directory name"

#===========================================================
# Check "filter"
#===========================================================

if [ -z "${filter}" ]; then
  filter=on
elif [ _"${filter}" = _"on" ] || [ _"${filter}" = _"off" ]; then
  :
else
  error_exit "${filter}: invalid filter name (on/off)"
fi

#===========================================================
# Define threads
#===========================================================

{
  unset max_threads tmp_threads
  max_threads=$(getconf _NPROCESSORS_ONLN)
  [ -z "$max_threads" ] && max_threads=$(getconf NPROCESSORS_ONLN)
  [ -z "$max_threads" ] && max_threads=$(ksh -c 'getconf NPROCESSORS_ONLN')
  [ -z "$max_threads" ] && max_threads=1
  tmp_threads=$(("${threads:-1}" + 0))
} 2>/dev/null || true

if [ "${tmp_threads:-0}" -gt 1 ] && [ "${tmp_threads}" -lt "${max_threads}" ]; then
  :
else
  threads=$(echo "${max_threads}" | awk '{print int($0*2/3+0.5)}')
fi

################################################################################
# Setting Conda environment
################################################################################

: >.DAJIN_temp/log.txt

cat <<EOF >&2
--------------------------------------------------------------------------------
Creat Conda environment
--------------------------------------------------------------------------------
EOF

. DAJIN/src/conda_setting.sh

################################################################################
# Formatting environments
################################################################################

#===========================================================
# Make temporal directory
#===========================================================

dirs="fasta fasta_conv fasta_ont NanoSim data"
echo "${dirs}" |
  sed "s:^:.DAJIN_temp/:g" |
  sed "s: : .DAJIN_temp/:g" |
  xargs mkdir -p

./DAJIN/src/format_fasta.sh "$design" "$input_dir" "$grna"

#===========================================================
# Define target mutation type (deletion:D, insertion:I, substitution:S)
#===========================================================

minimap2 -ax splice .DAJIN_temp/fasta/wt.fa .DAJIN_temp/fasta/target.fa --cs 2>/dev/null |
  grep -v "^@" |
  awk '{
    cstag=$(NF-1)
    if(cstag ~ "~") print "D"
    else if(cstag ~ /\+/) print "I"
    else if(cstag ~ /\*/) print "S"
    else if(cstag ~ /\-/) print "D"
    else print "D"
    }' >.DAJIN_temp/target_mutation_type

################################################################################
# NanoSim
################################################################################

cat <<EOF >&2
--------------------------------------------------------------------------------
NanoSim read simulation
--------------------------------------------------------------------------------
EOF

set +u
conda activate DAJIN_nanosim
set -u

if ! find .DAJIN_temp/fasta_ont/*simulated* >/dev/null 2>&1; then
  ./DAJIN/src/nanosim.sh "${control}" "${threads:-1}"
fi

################################################################################
# MIDS conversion
################################################################################

cat <<EOF >&2
--------------------------------------------------------------------------------
Preprocessing
--------------------------------------------------------------------------------
EOF

set +u
conda activate DAJIN
set -u

#===========================================================
# Get mutation loci
#===========================================================

minimap2 -ax splice .DAJIN_temp/fasta_conv/wt.fa .DAJIN_temp/fasta_conv/target.fa --cs 2>/dev/null |
  awk '{for(i=1; i<=NF;i++) if($i ~ /cs:Z/) print $i}' |
  sed -e "s/cs:Z:://g" -e "s/:/ /g" -e "s/~/ /g" |
  tr -d "\~\*\-\+atgc" |
  awk '{$NF=0; for(i=1;i<=NF;i++) sum+=$i} END{print $1,sum}' |
  cat >.DAJIN_temp/data/mutation_points

#===========================================================
# MIDS conversion
#===========================================================

find .DAJIN_temp/fasta_ont -type f |
  sort |
  awk '{print "./DAJIN/src/mids_classification.sh", $0, "wt", "&"}' |
  awk -v th="${threads:-1}" 'NR%th==0 {sub("&", "&\nwait")}1; END {print "wait"}' |
  sh - 2>/dev/null

################################################################################
# Prediction
################################################################################

cat <<EOF >&2
--------------------------------------------------------------------------------
Predict allele types
--------------------------------------------------------------------------------
EOF

./DAJIN/src/ml_prediction.sh "${control}" "${threads:-1}" \
  >.DAJIN_temp/data/DAJIN_MIDS_prediction_result.txt ||
  exit 1

###############################################################################
# Clustering
###############################################################################

cat <<EOF >&2
--------------------------------------------------------------------------------
Clustering alleles
--------------------------------------------------------------------------------
EOF

rm -rf .DAJIN_temp/clustering 2>/dev/null || true
mkdir -p .DAJIN_temp/clustering/temp

#===========================================================
# Prepare control's score to define sequencing error
#===========================================================

./DAJIN/src/clustering_control_scoring.sh "${control}" "${threads:-1}"

#===========================================================
# Clustering
#===========================================================

cat .DAJIN_temp/data/DAJIN_MIDS_prediction_result.txt |
  cut -f 2,3 |
  sort -u |
  awk -v ctrl="$control" '$1 $2 != ctrl "wt"' |
  awk -v th="${threads:-1}" '{print "./DAJIN/src/clustering.sh", $1, $2, th}' |
  sh -

cat .DAJIN_temp/data/DAJIN_MIDS_prediction_result.txt |
  awk -v ctrl="$control" '$2 $3 == ctrl "wt" {print $1"\t"1}' |
  cat >".DAJIN_temp/clustering/temp/hdbscan_${control}_wt".csv
cat ".DAJIN_temp/clustering/temp/hdbscan_${control}_wt".csv >.DAJIN_temp/clustering/temp/query_seq_${control}_wt.csv
true >".DAJIN_temp/clustering/temp/possible_true_mut_${control}_wt".csv

#===========================================================
# Allele percentage
#===========================================================

rm -rf ".DAJIN_temp/clustering/allele_per/" 2>/dev/null
mkdir -p ".DAJIN_temp/clustering/allele_per/"

cat .DAJIN_temp/data/DAJIN_MIDS_prediction_result.txt |
  cut -f 2 |
  sort -u |
  awk -v filter="${filter:-on}" '{print "./DAJIN/src/clustering_allele_percentage.sh", $1, filter, "&"}' |
  awk -v th="${threads:-1}" 'NR%th==0 {sub("&", "&\nwait")}1;END{if(NR!=th) print "wait"}' |
  sh -

################################################################################
# Get consensus sequence in each cluster
################################################################################

cat <<EOF >&2
--------------------------------------------------------------------------------
Report consensus sequence
--------------------------------------------------------------------------------
EOF

#===========================================================
# Setting directory
#===========================================================

rm -rf .DAJIN_temp/consensus 2>/dev/null || true
mkdir -p .DAJIN_temp/consensus/temp .DAJIN_temp/consensus/sam

#===========================================================
# Generate temporal SAM files
#===========================================================

cat .DAJIN_temp/clustering/allele_per/label* |
  cut -d " " -f 1,2 |
  sort -u |
  grep -v abnormal |
  while read -r barcode mapping_alleletype; do
    mapping_alleletype="${mapping_alleletype%.csv}"

    cat .DAJIN_temp/clustering/allele_per/readid_cl_mids_"${barcode}"_"${mapping_alleletype}".csv |
      awk '{print ">"$1}' |
      sort |
      cat >.DAJIN_temp/consensus/tmp_id

    cat .DAJIN_temp/fasta_ont/"${barcode}".fa |
      awk 'BEGIN{RS=">"}{print ">"$1" "$NF}' |
      sed 1d |
      sort |
      join - .DAJIN_temp/consensus/tmp_id |
      awk '{gsub(" ", "\n")}1' |
      minimap2 -ax map-ont -t "${threads:-1}" ".DAJIN_temp/fasta/${mapping_alleletype}.fa" - --cs=long 2>/dev/null |
      sort |
      cat >.DAJIN_temp/consensus/sam/"${barcode}"_"${mapping_alleletype}".sam
    rm .DAJIN_temp/consensus/tmp_id
  done

#===========================================================
# Execute consensus.sh
#===========================================================

rm -rf .DAJIN_temp/consensus/FASTA .DAJIN_temp/consensus/HTML 2>/dev/null || true
mkdir -p .DAJIN_temp/consensus/FASTA .DAJIN_temp/consensus/HTML

cat .DAJIN_temp/clustering/allele_per/label* |
  awk '{nr[$1]++; print $0, nr[$1]}' |
  grep -v abnormal |
  awk '{print "./DAJIN/src/consensus.sh", $0, "&"}' |
  awk -v th="${threads:-1}" 'NR%th==0 {sub("&", "&\nwait")}1;END{if(NR!=th) print "wait"}' |
  sh -

################################################################################
# Summarize to Details.csv and Details.pdf
################################################################################

./DAJIN/src/details.sh

################################################################################
# Mapping by minimap2 for IGV visualization
################################################################################

cat <<EOF >&2
--------------------------------------------------------------------------------
Generate BAM files
--------------------------------------------------------------------------------
EOF

if [ "$genome_check" -eq 0 ]; then
  ./DAJIN/src/generate_bam.sh "no_ref" "${threads:-1}"
else
  ./DAJIN/src/generate_bam.sh "${genome}" "${threads:-1}"
fi

################################################################################
# Move output files
################################################################################

rm -rf "${output_dir:=DAJIN_results}" 2>/dev/null || true
mkdir -p "${output_dir:-DAJIN_results}"/BAM
mkdir -p "${output_dir:-DAJIN_results}"/Consensus

#===========================================================
# BAM
#===========================================================

rm -rf .DAJIN_temp/bam/temp 2>/dev/null || true
cp -r .DAJIN_temp/bam/* "${output_dir:-DAJIN_results}"/BAM/ 2>/dev/null

#===========================================================
# Consensus
#===========================================================

(
  find .DAJIN_temp/consensus/* -type d |
    grep -v -e "consensus/temp" -e "sam" |
    xargs -I @ cp -f -r @ "${output_dir:-DAJIN_results}"/Consensus/
) 2>/dev/null || true

#===========================================================
# Details
#===========================================================

cp .DAJIN_temp/details/* "${output_dir:-DAJIN_results}"/

################################################################################
# Finish call
################################################################################

[ -z "${TEST}" ] && rm -rf .DAJIN_temp/

set +u
conda deactivate

cat <<EOF >&2
--------------------------------------------------------------------------------
Completed!
Check ${output_dir:-DAJIN_results} directory
--------------------------------------------------------------------------------
EOF

exit 0