#!/usr/bin/env bash
# Run on the server after upload:
#   bash run_trufor_remote.sh [gpu-id]
set -Eeuo pipefail

PROJECT_DIR=""
for candidate in \
  /data/users/summer2026/trufor/TruFor_train_test \
  /data/users/summer2026/TruFor/TruFor_train_test; do
  if [[ -f "$candidate/train.py" ]]; then
    PROJECT_DIR="$candidate"
    break
  fi
done
GPU_ID="${1:-0}"
ENV_NAME=trufor_cuda102

test -n "$PROJECT_DIR" || { echo "TruFor project directory was not found under /data/users/summer2026." >&2; exit 1; }

# Server driver: NVIDIA 440.xx / CUDA 10.2. The repository YAML instead pins CUDA 11.3.
source /opt/miniconda3/etc/profile.d/conda.sh
if ! conda env list | awk '{print $1}' | grep -qx "$ENV_NAME"; then
  conda create -y -n "$ENV_NAME" python=3.7 pytorch=1.11.0 torchvision=0.12.0 torchaudio=0.11.0 cudatoolkit=10.2 -c pytorch -c conda-forge
fi

# jpegio is intentionally omitted: TruFor does not import it, and it needs
# unavailable system headers (jpeglib.h) to build from source on this server.
# Keep these outside the environment-creation block so a failed first setup
# can be repaired simply by re-running this launcher.
conda run -n "$ENV_NAME" python -m pip install numpy==1.21.5 pillow==9.0.1 tqdm==4.64.0 scipy==1.7.3 scikit-image==0.16.2 pandas==1.3.5 pyyaml==6.0 albumentations==1.2.1 yacs==0.1.8 tensorboardX==2.5.1 mmcls==0.23.1 mmsegmentation==0.25.0 timm==0.5.4 opencv-python==4.5.5.64 torch-dct==0.1.5
# Do not leave both OpenCV distributions installed: they each provide cv2.
conda run -n "$ENV_NAME" python -m pip uninstall -y opencv-python-headless opencv-contrib-python
conda run -n "$ENV_NAME" python -m pip install --force-reinstall --no-deps opencv-python==4.5.5.64
conda run -n "$ENV_NAME" python -m pip install mmcv-full==1.5.3 -f https://download.openmmlab.com/mmcv/dist/cu102/torch1.11.0/index.html

cd "$PROJECT_DIR"
exec conda run --no-capture-output -n "$ENV_NAME" python train.py -g "$GPU_ID" -exp trufor_ph2