#!/bin/bash
# Обучение на TFRecord-шардах i1 — одной командой.
#
#   bash run.sh /tmp/hf/tst 10                       # 10 эпох по всем шардам в папке
#   bash run.sh /tmp/hf/tst 10 --init /path/to/step_061000.pt
#   bash run.sh /tmp/hf/tst 10 --step 61000 --batch 8 --fg
#
# Аргументы можно передавать и флагами, и через env (флаг приоритетнее):
#   --init PATH   стартовые веса: папка transformer/ или .pt-чекпойнт  (env INIT)
#   --step N      стартовый шаг счётчика и расписания lr               (env STEP)
#   --batch N     батч                                                  (env BATCH)
#   --out DIR     куда писать чекпойнты/лог                             (env OUT)
#   --epochs N    бюджет в эпохах (второй позиционный аргумент)         (env EPOCHS)
#   --fg          работать в переднем плане
# Для .pt-чекпойнта шаг берётся из файла, --step не нужен.
set -euo pipefail
cd "$(dirname "$0")"

DATA=""; EPOCHS="${EPOCHS:-10}"
INIT="${INIT:-transformer}"; STEP="${STEP:-61000}"
OUT="${OUT:-tmp/run}"; BATCH="${BATCH:-16}"; LOG="${LOG:-train.log}"; FG=0

while [ $# -gt 0 ]; do
  case "$1" in
    --init)   INIT="$2"; shift 2 ;;
    --step)   STEP="$2"; shift 2 ;;
    --batch)  BATCH="$2"; shift 2 ;;
    --out)    OUT="$2"; shift 2 ;;
    --epochs) EPOCHS="$2"; shift 2 ;;
    --log)    LOG="$2"; shift 2 ;;
    --fg)     FG=1; shift ;;
    -*)       echo "неизвестный флаг: $1"; exit 1 ;;
    *)        if [ -z "$DATA" ]; then DATA="$1"; else EPOCHS="$1"; fi; shift ;;
  esac
done

[ -n "$DATA" ] || { echo "укажи папку с TFRecord-шардами, например: bash run.sh /tmp/hf/tst 10"; exit 1; }
[ -e "$INIT" ] || { echo "нет стартовых весов: $INIT"; exit 1; }
[ -d "$DATA" ] || { echo "нет папки с шардами: $DATA"; exit 1; }
mkdir -p "$OUT"

STEP_ARGS=(--step "$STEP")
if [ -f "$INIT" ]; then                     # .pt помнит свой шаг — не перебиваем
  STEP_ARGS=()
  echo "стартовые веса:   $INIT (.pt, шаг из файла)"
else
  echo "стартовые веса:   $INIT (шаг $STEP)"
fi
echo "shard'ов в папке: $(ls -1 "$DATA"/*.tfrecord-* 2>/dev/null | wc -l)"
echo "бюджет:           $EPOCHS эпох, батч $BATCH, out $OUT"
echo "свободно на диске: $(df -h "$OUT" | awk 'NR==2 {print $4}')"

CMD=(python3 src/train.py
  --init "$INIT" "${STEP_ARGS[@]}"
  --tfrecord "$DATA" --epochs "$EPOCHS"
  --batch "$BATCH" --size 256 --text-len 192 --prefetch 256 --workers 8
  --lr 1e-4 --lr-min 1e-5 --lr-warmup 0 --ema 0.999
  --sample-every 2000 --save-every 4000 --log-every 200
  --out "$OUT")

echo "команда:          ${CMD[*]}"
echo
if [ "$FG" = "1" ]; then
  "${CMD[@]}"
else
  setsid nohup "${CMD[@]}" < /dev/null > "$LOG" 2>&1 &
  sleep 20
  echo "запущено в фоне. смотреть:  tail -f $LOG"
fi