#!/usr/bin/env bash

set -euo pipefail

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

if ip -o -4 addr show 2>/dev/null | awk '{split($4, address, "/"); print address[1]}' | grep -qx "$CURRENT_SERVER_IP"; then
  export OMNIGEN_CODE_ROOT="${OMNIGEN_CODE_ROOT:-/home/wenting/zr/gen_code}"
  DEFAULT_JSONL_PATH="$PROJECT_ROOT/dataset/cxr_radiomics_current_server/test_metadata.jsonl"
else
  export OMNIGEN_CODE_ROOT="${OMNIGEN_CODE_ROOT:-/NAS_REMOTE/vicky/wt/codes/cxr_syn/gen_code}"
  DEFAULT_JSONL_PATH="$PROJECT_ROOT/dataset/cxr_radiomics/test_metadata.jsonl"
fi

export PYTHONPATH="$PROJECT_ROOT:$OMNIGEN_CODE_ROOT:${PYTHONPATH:-}"
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
export HF_HOME="${HF_HOME:-/NAS_REMOTE/vicky/wt/huggingface/models}"
export HF_HUB_CACHE="${HF_HUB_CACHE:-/tmp/flow_grpo_hf_cache/hub}"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}"

MODEL_PATH="${MODEL_PATH:-Shitao/OmniGen-v1}"
LORA_PATH="${LORA_PATH:-$PROJECT_ROOT/logs/radiomics/omnigen-lora-r32-a64-4gpu-bs32-evalbs24-kl-beta0p005-from-cxr-lora30000/checkpoints/checkpoint-60/lora}"
JSONL_PATH="${JSONL_PATH:-$DEFAULT_JSONL_PATH}"
OUTPUT_DIR="${OUTPUT_DIR:-$PROJECT_ROOT/outputs/omnigen_lora_4gpu_ckp_60_eval}"
BATCH_SIZE="${BATCH_SIZE:-4}"
NUM_GPUS="${NUM_GPUS:-4}"
INFERENCE_STEPS="${INFERENCE_STEPS:-50}"
GUIDANCE_SCALE="${GUIDANCE_SCALE:-2.5}"
IMG_GUIDANCE_SCALE="${IMG_GUIDANCE_SCALE:-2.0}"
SEED="${SEED:-42}"

python3 "$PROJECT_ROOT/scripts/test_omnigen_cxr.py" \
  --model_path "$MODEL_PATH" \
  --lora_path "$LORA_PATH" \
  --omnigen_code_root "$OMNIGEN_CODE_ROOT" \
  --jsonl_path "$JSONL_PATH" \
  --output_dir "$OUTPUT_DIR" \
  --batch_size "$BATCH_SIZE" \
  --num_gpus "$NUM_GPUS" \
  --inference_steps "$INFERENCE_STEPS" \
  --guidance_scale "$GUIDANCE_SCALE" \
  --img_guidance_scale "$IMG_GUIDANCE_SCALE" \
  --seed "$SEED" \
  "$@"