#!/usr/bin/env bash
# Greedy one-shot linear clustering with mmseqs lin8clust (MMseqs2 fork).
# Build with -DMMSEQS_INT64_IDS=1 (required above ~4e9 sequences).
#
# Nodes coordinate through markers on a SHARED filesystem — no MPI, no SLURM.
# Run the SAME command on every machine, each with its own --node-id; the pipeline
# fans out with --threads and the passes meet at markers under $WORK. Rerunning
# resumes: each pass skips what it already finished.
#
#   single node :  ./run_oneshot_logan.sh single
#   many nodes  :  ./run_oneshot_logan.sh <node-id> <node-list>       (0-based id)
#                    on machine A:  ./run_oneshot_logan.sh 0 nodeA,nodeB
#                    on machine B:  ./run_oneshot_logan.sh 1 nodeA,nodeB
#                  ($INPUT and $WORK must be the same shared path on every machine)

set -euo pipefail

# ── paths ─────────────────────────────────────────────────────────────────────
MMSEQS=$HOME/MMseqs2/build64/src/mmseqs
INPUT=/fast2/imgyuri/LOGAN_TEST/nonhuman-complete.first.fa   # directory of *.fa.zst shards
WORK=/fast2/imgyuri/LOGAN_TEST/oneshot_work
RESULT=/fast2/imgyuri/LOGAN_TEST/oneshot_result/clu

# ── clustering parameters (identical for single and multi node) ────────────────
PAR="--min-seq-id 0.5 -c 0.9 --cov-mode 0 \
--kmer-per-seq 50 --num-adjacency 3 --clust-hash 1 \
--remove-tmp-files 1 -v 3"

mkdir -p "$WORK" "$(dirname "$RESULT")"
compgen -G "$INPUT/*.fa.zst" >/dev/null || { echo "no *.fa.zst under $INPUT" >&2; exit 1; }

# ── 1) single node ─────────────────────────────────────────────────────────────
# One machine: the aligning pass decides each group inline (fused path).
if [ "${1:-}" = "single" ]; then
    # shellcheck disable=SC2086
    "$MMSEQS" lin8clust "$INPUT"/*.fa.zst "$RESULT" "$WORK/tmp" \
        --threads "$(nproc)" $PAR
    echo "done: $RESULT ($RESULT.tsv)"
    exit 0
fi

# ── 2) multi node: run this on every machine, each with its own node-id ─────────
#   $1 = this machine's node-id (0-based)      $2 = comma list of all machines
NODE=${1:?usage: single  |  <node-id> <node-list>}
NODELIST=${2:?usage: <node-id> <node-list>, e.g. 0 nodeA,nodeB}
# shellcheck disable=SC2086
"$MMSEQS" lin8clust "$INPUT"/*.fa.zst "$RESULT" "$WORK/tmp" \
    --node-list "$NODELIST" --node-id "$NODE" \
    --threads "$(nproc)" $PAR
[ "$NODE" -eq 0 ] && echo "done: $RESULT ($RESULT.tsv)"