#!/usr/bin/env bash
# SUM Parts - full training run on the real dataset
#
# Usage:
#   bash train_full.sh                       # pointnet, paper settings
#   bash train_full.sh pointnext-xl          # auto-picks a voxel_max that fits
#   CFG_VOXEL_MAX=40000 bash train_full.sh pointnext-xl
#   EPOCHS=20 bash train_full.sh pointnet
#
# Measured on this box (RTX 3060 12GB), 24 train tiles, loop 30, batch_size 2
# => 360 iter/epoch:
#
#   pointnet        voxel_max 64000 (paper)   6.01G  0.291 s/iter   ~2.9 h/100ep
#   pointnet++msg   voxel_max 64000 (paper)   4.16G  0.675 s/iter   ~6.8 h/100ep
#   pointvector-xl  voxel_max 24000            6.46G  0.402 s/iter   ~4.0 h/100ep
#   pointnext-xl    voxel_max 32000            8.03G  0.635 s/iter   ~6.4 h/100ep
#
# The XL rows are NOT the paper configuration. At voxel_max 64000 they need
# ~15-16 GB, and on WSL2 the driver spills past VRAM into host RAM instead of
# raising OOM -- the run completes at 47 s/iter, i.e. ~20 days for 100 epochs.
# Numbers from a reduced voxel_max are not comparable to the published ones.
set -euo pipefail

CONDA_ROOT="$HOME/miniconda3"
ENV_NAME="sumparts"
SEG="$HOME/sum-parts/semantic_segmentation/PointNeXt_bundle/examples/segmentation"
DATA="$HOME/sum-parts/data/face_labeling/texsp_pcl"   # triangle track, 13 classes

CFG="${1:-pointnet}"
EPOCHS="${EPOCHS:-100}"
VAL_FREQ="${VAL_FREQ:-5}"       # cfg default is 1; every epoch costs 8 val tiles
LOG="${LOG:-/tmp/sumparts_train_${CFG}.log}"

# largest voxel_max measured to stay inside 12 GB for each model
if [ -n "${CFG_VOXEL_MAX:-}" ]; then
    VOXEL_MAX="$CFG_VOXEL_MAX"
else
    case "$CFG" in
        pointnext-xl)   VOXEL_MAX=32000 ;;
        pointvector-xl) VOXEL_MAX=24000 ;;
        *)              VOXEL_MAX=64000 ;;   # paper setting; small models fit
    esac
fi

source "$CONDA_ROOT/etc/profile.d/conda.sh"
conda activate "$ENV_NAME"
export WANDB_MODE=disabled WANDB_SILENT=true CUDA_HOME="$CONDA_PREFIX"

[ -d "$DATA/train" ] || { echo "error: $DATA/train missing -- run download_data.sh and prepare_full_split.sh" >&2; exit 1; }

echo "=== full training ==="
echo "  cfg        : $CFG"
echo "  data_root  : $DATA"
echo "  train/val/test : $(find -L "$DATA/train" -name '*.ply' | wc -l)/$(find -L "$DATA/val" -name '*.ply' | wc -l)/$(find -L "$DATA/test" -name '*.ply' | wc -l)"
echo "  epochs     : $EPOCHS"
echo "  val_freq   : $VAL_FREQ"
echo "  voxel_max  : $VOXEL_MAX$([ "$VOXEL_MAX" -lt 64000 ] && echo '  (reduced from paper 64000 to fit VRAM)')"
echo "  log        : $LOG"
echo

cd "$SEG"

set +e
python -u main.py \
    --cfg "../../cfgs/sumv2_triangle/${CFG}.yaml" \
    mode=train \
    dataset.common.data_root="$DATA" \
    dataset.train.voxel_max="$VOXEL_MAX" \
    epochs="$EPOCHS" \
    val_freq="$VAL_FREQ" \
    wandb.use_wandb=False \
    > "$LOG" 2>&1
rc=$?
set -e

echo "=== exit=$rc | last 25 lines ==="
tail -25 "$LOG"
[ "$rc" -eq 0 ] && echo "TRAIN DONE" || echo "TRAIN FAILED (rc=$rc)"
exit "$rc"