#!/usr/bin/env bash
set -euo pipefail

# 4 GPU radiomics OmniGen LoRA training.
# Override config/port if needed:
# CONFIG_ENTRY=config/grpo.py:general_radiomics_omnigen_4gpu_kl MAIN_PROCESS_PORT=29502 bash scripts/single_node/grpo_4gpu.sh

PROJECT_ROOT=$(CDPATH= cd -- "$(dirname -- "$0")/../.." && pwd)
export PROJECT_ROOT

WANDB_ENV_FILE="${WANDB_ENV_FILE:-$HOME/.wandb_env}"
if [[ -f "$WANDB_ENV_FILE" ]]; then
  # shellcheck disable=SC1090
  source "$WANDB_ENV_FILE"
fi

export OMNIGEN_CODE_ROOT=${OMNIGEN_CODE_ROOT:-/home/wenting/gen_joint}
if [ ! -f "$OMNIGEN_CODE_ROOT/OmniGen/__init__.py" ]; then
  echo "Missing OmniGen package at $OMNIGEN_CODE_ROOT/OmniGen/__init__.py" >&2
  exit 1
fi

export PYTHONPATH="$OMNIGEN_CODE_ROOT:$PROJECT_ROOT:${PYTHONPATH:-}"
export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}
export HF_HOME="${HF_HOME:-/home/wenting/.cache/huggingface}"
export HF_HUB_CACHE="${HF_HUB_CACHE:-/home/wenting/.cache/huggingface/hub}"
mkdir -p "$HF_HUB_CACHE"
export WANDB_MODE="${WANDB_MODE:-online}"

if [[ "$WANDB_MODE" == "online" && -z "${WANDB_API_KEY:-}" ]]; then
  echo "[ERROR] WANDB_MODE=online but WANDB_API_KEY is not set." >&2
  echo "Please export WANDB_API_KEY or put it in $WANDB_ENV_FILE." >&2
  exit 1
fi

export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}
export RADIOMICS_REWARD_FAIL_FAST=${RADIOMICS_REWARD_FAIL_FAST:-1}

export SAVE_FREQ="${SAVE_FREQ:-5}" 

if [[ "${SKIP_PREFLIGHT:-0}" != "1" ]]; then
  python3 "$PROJECT_ROOT/scripts/debug_preflight_omnigen_env.py"
fi

if [[ "${SKIP_REWARD_SMOKE:-0}" != "1" ]]; then
  python3 "$PROJECT_ROOT/scripts/debug_radiomics_reward_smoke.py" \
    --config "${CONFIG_ENTRY:-config/grpo.py:general_radiomics_omnigen_4gpu_kl}"
fi

python3 -m accelerate.commands.launch \
  --config_file scripts/accelerate_configs/multi_gpu.yaml \
  --num_processes=${NUM_PROCESSES:-4} \
  --main_process_port ${MAIN_PROCESS_PORT:-29501} \
  scripts/train_omnigen.py \
  --config "${CONFIG_ENTRY:-config/grpo.py:general_radiomics_omnigen_4gpu_kl}" \
  --config.save_freq="${SAVE_FREQ}"