#!/usr/bin/env bash
# SUM Parts - pipeline smoke test
#
# Goal is NOT accuracy. Goal is proving the whole chain survives:
#   ply load -> grid subsample -> cuda ops -> forward -> backward -> val
# Uses the lightest model (pointnet) and 1 epoch. Expect garbage mIoU.
#
# Usage:
#   bash smoke_train.sh                      # pointnet, 1 epoch
#   bash smoke_train.sh pointnext-xl 2       # other cfg / epoch count
set -euo pipefail

CONDA_ROOT="$HOME/miniconda3"
ENV_NAME="sumparts"
REPO="$HOME/sum-parts/semantic_segmentation/PointNeXt_bundle"

CFG="${1:-pointnet}"
EPOCHS="${2:-1}"

source "$CONDA_ROOT/etc/profile.d/conda.sh"
conda activate "$ENV_NAME"

export CUDA_HOME="$CONDA_PREFIX"
export PATH="$CUDA_HOME/bin:$PATH"

# main.py imports wandb at module scope, so the package must exist even with
# wandb.use_wandb=False. Disabled mode keeps it from asking for a login.
export WANDB_MODE=disabled
export WANDB_SILENT=true

cd "$REPO/examples/segmentation"

# Two constraints squeeze BS from both sides, and both fail confusingly:
#
#   too large -> the loader drops the last partial batch, so a single tile with
#     the cfg's batch_size (6 for pointnet) yields ZERO batches. The confusion
#     matrix is then never populated and cm.all_metrics() dies on
#         AttributeError: 'int' object has no attribute 'diag'
#
#   too small -> BatchNorm1d refuses a batch of 1 while training:
#         ValueError: Expected more than 1 value per channel when training,
#                     got input size torch.Size([1, 512])
#
# So: BS=2 (the BatchNorm floor), and LOOP repeats the tile enough times to fill
# whole batches.
LOOP="${3:-4}"
BS="${BS:-2}"

LOG="${LOG:-/tmp/sumparts_smoke.log}"

echo "=== smoke train: cfg=$CFG epochs=$EPOCHS loop=$LOOP bs=$BS ==="
echo "log: $LOG"

# Capture both streams to a file. Tracebacks were getting lost when this ran
# under a wrapper that only kept stdout, which made crashes look silent.
set +e
python -u main.py \
    --cfg "../../cfgs/sumv2_triangle/${CFG}.yaml" \
    mode=train \
    epochs="$EPOCHS" \
    wandb.use_wandb=False \
    dataset.train.loop="$LOOP" \
    batch_size="$BS" \
    val_batch_size=1 \
    val_freq=1 \
    > "$LOG" 2>&1
rc=$?
set -e

echo "=== exit=$rc | last 30 lines ==="
tail -30 "$LOG"

[ "$rc" -eq 0 ] && echo "SMOKE DONE" || echo "SMOKE FAILED (rc=$rc)"
exit "$rc"