#!/bin/bash
# Usage: stage_ds.sh <stage-name> <datasets-csv>
# Converts + loudness-norms + VQ-extracts one batch of datasets into
# /opt/work/data-fs (train) and /opt/work/data-fs-val (val).
set -euo pipefail
STAGE="$1"; DS="$2"
export PATH="$HOME/.local/bin:$PATH"
cd /opt/work/fish-speech
PY=/opt/work/fish-speech/.venv/bin/python

RAWD=/opt/work/staging/${STAGE}-raw
RAWV=/opt/work/staging/${STAGE}-raw-val
NORMD=/opt/work/staging/${STAGE}-norm
NORMV=/opt/work/staging/${STAGE}-norm-val

echo "== [$STAGE 1/5] prepare ($DS) =="
$PY /opt/work/scripts/prepare_data.py --datasets "$DS" \
    --out $RAWD --out-val $RAWV --holdout-suffix "_$STAGE"

echo "== [$STAGE 2/5] loudness norm =="
uvx --from fish-audio-preprocess fap loudness-norm $RAWD $NORMD --clean --num-workers 28
if [ -n "$(find $RAWV -name '*.wav' 2>/dev/null | head -1)" ]; then
  uvx --from fish-audio-preprocess fap loudness-norm $RAWV $NORMV --clean --num-workers 28
fi

echo "== [$STAGE 3/5] copy .lab =="
rsync -a --include='*/' --include='*.lab' --exclude='*' $RAWD/ $NORMD/
[ -d $NORMV ] && rsync -a --include='*/' --include='*.lab' --exclude='*' $RAWV/ $NORMV/ || true

echo "== [$STAGE 4/5] extract VQ (s2-pro codec) =="
$PY tools/vqgan/extract_vq.py $NORMD \
    --num-workers 1 --batch-size 12 \
    --config-name "modded_dac_vq" \
    --checkpoint-path "/opt/work/checkpoints/s2-pro/codec.pth"
if [ -d $NORMV ]; then
  $PY tools/vqgan/extract_vq.py $NORMV \
      --num-workers 1 --batch-size 12 \
      --config-name "modded_dac_vq" \
      --checkpoint-path "/opt/work/checkpoints/s2-pro/codec.pth"
fi

echo "== [$STAGE 5/5] merge into /opt/work/data-fs =="
mkdir -p /opt/work/data-fs /opt/work/data-fs-val
rsync -a $NORMD/ /opt/work/data-fs/
[ -d $NORMV ] && rsync -a $NORMV/ /opt/work/data-fs-val/ || true

echo "STAGE_${STAGE}_DONE"