#!/usr/bin/env bash
set -euo pipefail

# Created by argbash-init v2.10.0
# ARG_OPTIONAL_SINGLE([threads],[t],[Number of threads],[currently idle CPU cores])
# ARG_OPTIONAL_SINGLE([memory],[m],[Amount of memory in GB],[currently free memory])
# ARG_OPTIONAL_SINGLE([limit],[l],[Time limit],[no limit])
# ARG_OPTIONAL_SINGLE([out],[o],[Output directory],[./o])
# ARG_OPTIONAL_REPEATED([use],[u],[Programs to use (repeatable)])
# ARG_POSITIONAL_INF([files],[Files to go through],[1])
# ARG_OPTIONAL_SINGLE([glob],[],[What files to match in the directory],[*])
# ARGBASH_SET_DELIM([ =])
# ARG_OPTION_STACKING([getopt])
# ARG_RESTRICT_VALUES([no-local-options])
# ARG_DEFAULTS_POS([])
# ARG_HELP([<The general help message of my script>])
# ARGBASH_GO()
# needed because of Argbash --> m4_ignore([
### START OF CODE GENERATED BY Argbash v2.10.0 one line above ###
# Argbash is a bash code generator used to get arguments parsing right.
# Argbash is FREE SOFTWARE, see https://argbash.io for more info

die() {
  local _ret="${2:-1}"
  test "${_PRINT_HELP:-no}" = yes && print_help >&2
  echo "$1" >&2
  exit "${_ret}"
}

evaluate_strictness() {
  [[ "$2" =~ ^-(-(threads|memory|limit|out|use|files|glob|help)$|[tmlouh]) ]] && die "You have passed '$2' as a value of argument '$1', which makes it look like that you have omitted the actual value, since '$2' is an option accepted by this script. This is considered a fatal error."
}

begins_with_short_option() {
  local first_option all_short_options='tmlouh'
  first_option="${1:0:1}"
  test "$all_short_options" = "${all_short_options/$first_option/}" && return 1 || return 0
}

# THE DEFAULTS INITIALIZATION - POSITIONALS
_positionals=()
_arg_files=('')
# THE DEFAULTS INITIALIZATION - OPTIONALS
_arg_threads="currently idle CPU cores"
_arg_memory="currently free memory"
_arg_limit="no limit"
_arg_out="./o"
_arg_use=()
_arg_glob="*"

print_help() {
  printf '%s\n' "<The general help message of my script>"
  printf 'Usage: %s [-t|--threads <arg>] [-m|--memory <arg>] [-l|--limit <arg>] [-o|--out <arg>] [-u|--use <arg>] [--glob <arg>] [-h|--help] <files-1> [<files-2>] ... [<files-n>] ...\n' "$0"
  printf '\t%s\n' "<files>: Files to go through"
  printf '\t%s\n' "-t, --threads: Number of threads (default: 'currently idle CPU cores')"
  printf '\t%s\n' "-m, --memory: Amount of memory in GB (default: 'currently free memory')"
  printf '\t%s\n' "-l, --limit: Time limit (default: 'no limit')"
  printf '\t%s\n' "-o, --out: Output directory (default: './o')"
  printf '\t%s\n' "-u, --use: Programs to use (repeatable) (empty by default)"
  printf '\t%s\n' "--glob: What files to match in the directory (default: '*')"
  printf '\t%s\n' "-h, --help: Prints help"
}

parse_commandline() {
  _positionals_count=0
  while test $# -gt 0; do
    _key="$1"
    case "$_key" in
    -t | --threads)
      test $# -lt 2 && die "Missing value for the optional argument '$_key'." 1
      _arg_threads="$2"
      shift
      evaluate_strictness "$_key" "$_arg_threads"
      ;;
    --threads=*)
      _arg_threads="${_key##--threads=}"
      evaluate_strictness "$_key" "$_arg_threads"
      ;;
    -t*)
      _arg_threads="${_key##-t}"
      evaluate_strictness "$_key" "$_arg_threads"
      ;;
    -m | --memory)
      test $# -lt 2 && die "Missing value for the optional argument '$_key'." 1
      _arg_memory="$2"
      shift
      evaluate_strictness "$_key" "$_arg_memory"
      ;;
    --memory=*)
      _arg_memory="${_key##--memory=}"
      evaluate_strictness "$_key" "$_arg_memory"
      ;;
    -m*)
      _arg_memory="${_key##-m}"
      evaluate_strictness "$_key" "$_arg_memory"
      ;;
    -l | --limit)
      test $# -lt 2 && die "Missing value for the optional argument '$_key'." 1
      _arg_limit="$2"
      shift
      evaluate_strictness "$_key" "$_arg_limit"
      ;;
    --limit=*)
      _arg_limit="${_key##--limit=}"
      evaluate_strictness "$_key" "$_arg_limit"
      ;;
    -l*)
      _arg_limit="${_key##-l}"
      evaluate_strictness "$_key" "$_arg_limit"
      ;;
    -o | --out)
      test $# -lt 2 && die "Missing value for the optional argument '$_key'." 1
      _arg_out="$2"
      shift
      evaluate_strictness "$_key" "$_arg_out"
      ;;
    --out=*)
      _arg_out="${_key##--out=}"
      evaluate_strictness "$_key" "$_arg_out"
      ;;
    -o*)
      _arg_out="${_key##-o}"
      evaluate_strictness "$_key" "$_arg_out"
      ;;
    -u | --use)
      test $# -lt 2 && die "Missing value for the optional argument '$_key'." 1
      _arg_use+=("$2")
      shift
      evaluate_strictness "$_key" "${_arg_use[-1]}"
      ;;
    --use=*)
      _arg_use+=("${_key##--use=}")
      evaluate_strictness "$_key" "${_arg_use[-1]}"
      ;;
    -u*)
      _arg_use+=("${_key##-u}")
      evaluate_strictness "$_key" "${_arg_use[-1]}"
      ;;
    --glob)
      test $# -lt 2 && die "Missing value for the optional argument '$_key'." 1
      _arg_glob="$2"
      shift
      evaluate_strictness "$_key" "$_arg_glob"
      ;;
    --glob=*)
      _arg_glob="${_key##--glob=}"
      evaluate_strictness "$_key" "$_arg_glob"
      ;;
    -h | --help)
      print_help
      exit 0
      ;;
    -h*)
      print_help
      exit 0
      ;;
    *)
      _last_positional="$1"
      _positionals+=("$_last_positional")
      _positionals_count=$((_positionals_count + 1))
      ;;
    esac
    shift
  done
}

handle_passed_args_count() {
  local _required_args_string="'files'"
  test "${_positionals_count}" -ge 1 || _PRINT_HELP=yes die "FATAL ERROR: Not enough positional arguments - we require at least 1 (namely: $_required_args_string), but got only ${_positionals_count}." 1
}

assign_positional_args() {
  local _positional_name _shift_for=$1
  _positional_names="_arg_files "
  _our_args=$((${#_positionals[@]} - 1))
  for ((ii = 0; ii < _our_args; ii++)); do
    _positional_names="$_positional_names _arg_files[$((ii + 1))]"
  done

  shift "$_shift_for"
  for _positional_name in ${_positional_names}; do
    test $# -gt 0 || break
    eval "$_positional_name=\${1}" || die "Error during argument parsing, possibly an Argbash bug." 1
    shift
  done
}

parse_commandline "$@"
handle_passed_args_count
assign_positional_args 1 "${_positionals[@]}"

# OTHER STUFF GENERATED BY Argbash

### END OF CODE GENERATED BY Argbash (sortof) ### ])
# [ <-- needed because of Argbash

# vvv  PLACE YOUR CODE HERE  vvv

: "${_POLAP_DEBUG:=0}"
export _POLAP_DEBUG
: "${_POLAP_RELEASE:=0}"
export _POLAP_RELEASE

get_free_cpus() {
  local total_cpus used_cpus load_avg free_cpus

  total_cpus=$(nproc)
  # Use the 1-minute load average and round up
  load_avg=$(awk '{print int($1 + 0.999)}' </proc/loadavg)
  used_cpus=$((load_avg > total_cpus ? total_cpus : load_avg))
  free_cpus=$((total_cpus - used_cpus))

  echo "$free_cpus"
}

get_available_memory_gb() {
  awk '/^MemAvailable:/ {printf "%.0f\n", $2 / 1024 / 1024}' /proc/meminfo
}

# Parse positional arguments
_input_files=("${_arg_files[@]:0:${#_arg_files[@]}-1}")
_output_file="${_arg_files[-1]}"

# Print parsed arguments
_arg_threads=$(get_free_cpus)
_arg_memory=$(get_available_memory_gb)
echo "Threads: $_arg_threads"
echo "Memory: $_arg_memory GB"
echo "Limit: $_arg_limit"
echo "Output dir: $_arg_out"
echo "Use programs: ${_arg_use[@]}"
echo "Input files: ${_input_files[@]}"
echo "Output file: $_output_file"

source "$(conda info --base)/etc/profile.d/conda.sh"
if [[ "$CONDA_DEFAULT_ENV" != "bolap" ]]; then
  echo "You're not in the bolap environment. Chaniging 'bolap'..."
  conda activate bolap
fi

# Check for seqkit
if ! command -v seqkit &>/dev/null; then
  echo "ERROR: seqkit is required for this script." >&2
  exit 1
fi

trap 'echo "❌ Error at line $LINENO: Command \"${BASH_COMMAND}\" failed." >&2' ERR

log() {
  echo "[$(date +'%Y-%m-%d %H:%M:%S')] $*" >&2
}

get_read_cmd() {
  [[ "$1" == *.gz ]] && echo "zcat" || echo "cat"
}

estimate_len_and_q_ends() {
  local fq="$1"
  local read_cmd
  read_cmd=$(get_read_cmd "$fq")

  local total_len=0
  local total_q=0
  local line_count=0

  # Read 10,000 lines from start and 10,000 from end
  {
    $read_cmd "$fq" | head -n 100
    # $read_cmd "$fq" | tail -n 100
  } | while read -r line; do
    line_count=$((line_count + 1))
    mod=$((line_count % 4))
    case $mod in
    2) seq_line="$line" ;;
    0)
      qual_line="$line"
      total_len=$((total_len + ${#seq_line}))
      for ((i = 0; i < ${#qual_line}; i++)); do
        qscore=$(printf '%d' "'${qual_line:$i:1}")
        total_q=$((total_q + qscore - 33))
      done
      ;;
    esac
  done

  local reads=$((line_count / 4))
  [[ "$reads" -eq 0 || "$total_len" -eq 0 ]] && {
    echo "1 1"
    return
  }
  local avg_len=$((total_len / reads))
  local avg_q=$((total_q / total_len))
  echo "$avg_len $avg_q"
}

detect_seq_tech_by_length_q() {
  local fq="$1"
  read avg_len avg_q < <(estimate_len_and_q_ends "$fq")

  if [[ "$avg_len" -le 1000 ]]; then
    if [[ "$avg_q" -ge 30 ]]; then
      echo "illumina"
    else
      echo "short-read (low quality)"
    fi
  else
    if [[ "$avg_q" -ge 20 ]]; then
      echo "pacbio-hifi"
    elif [[ "$avg_len" -gt 10000 ]]; then
      if [[ "$avg_q" -ge 16 ]]; then
        echo "pacbio-corrected"
      else
        echo "pacbio-raw"
      fi
    elif [[ "$avg_q" -ge 12 ]]; then
      echo "nano-hq"
    else
      echo "nano-raw"
    fi
  fi
}

SAMPLE_READS=2500

get_seq_and_qual() {
  local fq="$1"
  local lines=$((SAMPLE_READS * 4))

  {
    zcat -f "$fq" | head -n "$lines"
    zcat -f "$fq" | tail -n "$lines"
  } | seqkit fx2tab -q -n -i 2>/dev/null
}

detect_seq_tech_seqkit() {
  local fq="$1"

  local total_len=0
  local total_q=0
  local total_bases=0
  local n_reads=0

  while IFS=$'\t' read -r _name seq qual; do
    len=${#seq}
    total_len=$((total_len + len))
    for ((i = 0; i < len; i++)); do
      q=$(($(printf "%d" "'${qual:$i:1}") - 33))
      total_q=$((total_q + q))
    done
    total_bases=$((total_bases + len))
    n_reads=$((n_reads + 1))
  done < <(get_seq_and_qual "$fq")

  [[ "$n_reads" -eq 0 || "$total_bases" -eq 0 ]] && {
    echo -e "$fq\tunknown\t0\t0"
    return
  }

  avg_len=$((total_len / n_reads))
  avg_q=$((total_q / total_bases))

  # Classification
  if [[ "$avg_len" -le 1000 ]]; then
    [[ "$avg_q" -ge 30 ]] && tech="illumina" || tech="short-read (low quality)"
  elif [[ "$avg_q" -ge 20 ]]; then
    tech="pacbio-hifi"
  elif [[ "$avg_len" -gt 10000 ]]; then
    [[ "$avg_q" -ge 16 ]] && tech="pacbio-corrected" || tech="pacbio-raw"
  elif [[ "$avg_q" -ge 12 ]]; then
    tech="nano-hq"
  else
    tech="nano-raw"
  fi

  echo -e "$fq\t$tech\t$avg_len\t$avg_q"
}

detect_seq_tech_seqkit_vars() {
  local fq="$1"
  local SAMPLE_READS=2500
  local lines=$((SAMPLE_READS * 4))

  local total_len=0
  local total_q="0.0"
  local n_reads=0

  while IFS=$'\t' read -r _id len qual; do
    total_len=$((total_len + len))
    total_q=$(awk -v tq="$total_q" -v q="$qual" -v l="$len" 'BEGIN { printf "%.6f", tq + (q * l) }')
    n_reads=$((n_reads + 1))
  done < <(
    {
      zcat -f "$fq" | head -n "$lines"
      zcat -f "$fq" | tail -n "$lines"
    } | seqkit fx2tab -n -i -l -q 2>/dev/null
  )

  [[ "$n_reads" -eq 0 || "$total_len" -eq 0 ]] && {
    seq_tech="unknown"
    avg_len=0
    avg_q=0
    return
  }

  avg_len=$((total_len / n_reads))
  avg_q=$(awk -v tq="$total_q" -v b="$total_len" 'BEGIN { printf "%.2f", tq / b }')

  # Classification logic
  if ((avg_len <= 500)); then
    if (($(awk "BEGIN {print ($avg_q >= 30)}"))); then
      seq_tech="illumina"
    else
      seq_tech="short-read (low quality)"
    fi
  elif (($(awk "BEGIN {print ($avg_q < 12)}"))); then
    seq_tech="nano-raw"
  elif (($(awk "BEGIN {print ($avg_q < 16)}"))); then
    if ((avg_len < 8000)); then
      seq_tech="pacbio-raw"
    else
      seq_tech="nano-corrected"
    fi
  elif (($(awk "BEGIN {print ($avg_q < 20)}"))); then
    if ((avg_len < 8000)); then
      seq_tech="pacbio-corrected"
    else
      seq_tech="nano-hq"
    fi
  else
    if ((avg_len > 10000)); then
      seq_tech="pacbio-hifi"
    else
      seq_tech="nano-hq"
    fi
  fi
}

# === MAIN ===

for fq in ${_input_files[@]}; do
  echo "Processing: $fq" >&2
  detect_seq_tech_seqkit_vars "$fq"

  echo "File: $fq"
  echo "Technology: $seq_tech"
  echo "Average Length: $avg_len"
  echo "Average Q: $avg_q"
done

echo "✅ Done. Report"

# | Technology           | Avg Read Length | Avg Q-score | Notes                  |
# | -------------------- | --------------- | ----------- | ---------------------- |
# | **Illumina**         | < 1000          | ≥ 30        | Short reads, high Q    |
# | **PacBio HiFi**      | > 1000          | ≥ 20        | Long, high-quality     |
# | **PacBio Corrected** | > 10000         | 16–19       | CLR with polishing     |
# | **PacBio Raw**       | > 10000         | < 16        | CLR uncorrected        |
# | **Nano HQ**          | > 1000          | 12–15       | Basecalled ONT         |
# | **Nano Raw**         | > 1000          | < 12        | Raw ONT, pre-polishing |

# ^^^  TERMINATE YOUR CODE BEFORE THE BOTTOM ARGBASH MARKER  ^^^

# ] <-- needed because of Argbash
