#!/usr/bin/env bash
# Full UVDoc training (no Doc3D), GPU 0, recommended hyperparameters.
# Stop with: kill <pid>  (or Ctrl+C if foreground)

set -euo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
PY="${PYTHON:-/root/miniconda3/envs/o3dedit/bin/python}"

UV_ROOT="${UV_ROOT:-$ROOT/UVDoc_final}"
LOGDIR="${LOGDIR:-$ROOT/log_full_uvdoc_gpu0}"

# Shorter uvdoc-only schedule (override with N_EPOCHS, N_DECAY env vars).
BS="${BS:-8}"
N_EPOCHS="${N_EPOCHS:-25}"
N_DECAY="${N_DECAY:-10}"
LR="${LR:-0.0002}"
EP_GAMMA="${EP_GAMMA:-10}"
VAL_RATIO="${VAL_RATIO:-0.05}"
SPLIT_SEED="${SPLIT_SEED:-42}"
NUM_WORKERS="${NUM_WORKERS:-8}"

exec "$PY" "$ROOT/UVDoc_official/train.py" \
  --data_to_use uvdoc \
  --data_path_UVDoc "$UV_ROOT" \
  --uvdoc_val_ratio "$VAL_RATIO" \
  --uvdoc_split_seed "$SPLIT_SEED" \
  --batch_size "$BS" \
  --n_epochs "$N_EPOCHS" \
  --n_epochs_decay "$N_DECAY" \
  --lr "$LR" \
  --alpha_w 5.0 \
  --beta_w 5.0 \
  --gamma_w 1.0 \
  --ep_gamma_start "$EP_GAMMA" \
  --appearance_augmentation visual noise color \
  --geometric_augmentationsUVDoc rotate \
  --num_workers "$NUM_WORKERS" \
  --device cuda:0 \
  --logdir "$LOGDIR"