#!/bin/bash -eu
set -o pipefail
{
#Shuffles and copies selfplay training from selfplay/ to shuffleddata/current/
#Should be run periodically.

if [[ $# -lt 4 ]]
then
    echo "Usage: $0 BASEDIR TMPDIR NTHREADS BATCHSIZE"
    echo "Currently expects to be run from within the 'python' directory of the KataGo repo, or otherwise in the same dir as shuffle.py."
    echo "BASEDIR containing selfplay data and models and related directories"
    echo "TMPDIR scratch space, ideally on fast local disk, unique to this loop"
    echo "NTHREADS number of parallel threads/processes to use in shuffle"
    echo "BATCHSIZE number of samples to concat together per batch for training"
    exit 0
fi
BASEDIR="$1"
shift
TMPDIR="$1"
shift
NTHREADS="$1"
shift
BATCHSIZE="$1"
shift

#------------------------------------------------------------------------------

OUTDIR=$(date "+%Y%m%d-%H%M%S")

mkdir -p "$BASEDIR"/shuffleddata/"$OUTDIR".tmp
mkdir -p "$TMPDIR"/train
mkdir -p "$TMPDIR"/val

echo "Beginning shuffle at" $(date "+%Y-%m-%d %H:%M:%S")

if [[ -n "${SKIP_VALIDATE:-}" ]]
then
  (
      time python3 ./shuffle.py \
           "$BASEDIR"/selfplay/ \
           -expand-window-per-row 0.4 \
           -taper-window-exponent 0.65 \
           -out-dir "$BASEDIR"/shuffleddata/"$OUTDIR".tmp/train \
           -out-tmp-dir "$TMPDIR"/train \
           -approx-rows-per-out-file 70000 \
           -num-processes "$NTHREADS" \
           -batch-size "$BATCHSIZE" \
           -only-include-md5-path-prop-lbound 0.00 \
           -only-include-md5-path-prop-ubound 1.00 \
           -output-npz \
           "$@" \
           2>&1 | tee "$BASEDIR"/shuffleddata/"$OUTDIR".tmp/outtrain.txt &

      wait
  )
else
  # Randomly peels off 5% of files generated by selfplay as validation data
  (
      time python3 ./shuffle.py \
           "$BASEDIR"/selfplay/ \
           -expand-window-per-row 0.4 \
           -taper-window-exponent 0.65 \
           -out-dir "$BASEDIR"/shuffleddata/"$OUTDIR".tmp/val \
           -out-tmp-dir "$TMPDIR"/val \
           -approx-rows-per-out-file 70000 \
           -num-processes "$NTHREADS" \
           -batch-size "$BATCHSIZE" \
           -only-include-md5-path-prop-lbound 0.95 \
           -only-include-md5-path-prop-ubound 1.00 \
           -output-npz \
           "$@" \
           2>&1 | tee "$BASEDIR"/shuffleddata/"$OUTDIR".tmp/outval.txt &

      wait
  )
  (
      time python3 ./shuffle.py \
           "$BASEDIR"/selfplay/ \
           -expand-window-per-row 0.4 \
           -taper-window-exponent 0.65 \
           -out-dir "$BASEDIR"/shuffleddata/"$OUTDIR".tmp/train \
           -out-tmp-dir "$TMPDIR"/train \
           -approx-rows-per-out-file 70000 \
           -num-processes "$NTHREADS" \
           -batch-size "$BATCHSIZE" \
           -only-include-md5-path-prop-lbound 0.00 \
           -only-include-md5-path-prop-ubound 0.95 \
           -output-npz \
           "$@" \
           2>&1 | tee "$BASEDIR"/shuffleddata/"$OUTDIR".tmp/outtrain.txt &

      wait
  )
fi
#set +x

#Just in case, give a little time for nfs
sleep 10

# Shuffle no longer maintains a "current" directory, instead training script searches for latest dir
# rm -f "$BASEDIR"/shuffleddata/current_tmp
# ln -s $OUTDIR "$BASEDIR"/shuffleddata/current_tmp
# mv -Tf "$BASEDIR"/shuffleddata/current_tmp "$BASEDIR"/shuffleddata/current

mv "$BASEDIR"/shuffleddata/"$OUTDIR".tmp "$BASEDIR"/shuffleddata/"$OUTDIR"

# CLEANUP ---------------------------------------------------------------

#Among shuffled dirs older than 2 hours, remove all but the most recent 5 of them.
#This should be VERY conservative and allow plenty of time for the training to switch
#to newer ones as they get generated.
echo "Cleaning up any old dirs"
find "$BASEDIR"/shuffleddata/ -mindepth 1 -maxdepth 1 -type d -mmin +120 | sort | head -n -5 | xargs --no-run-if-empty rm -r

echo "Finished shuffle at" $(date "+%Y-%m-%d %H:%M:%S")
#Make a little space between shuffles
echo ""
echo ""

exit 0
}
