#!/usr/bin/env bash

set -euo pipefail

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

source "$PROJECT_ROOT/scripts/single_node/server_profile_defaults.sh"
apply_server_profile_defaults

DEFAULT_BASE_DIR="/apdcephfs_zwfy2/share_303944931/connorxian/Code/medUni/wt/zr"
export OMNIGEN_CODE_ROOT="${OMNIGEN_CODE_ROOT:-${DEFAULT_BASE_DIR}/gen_joint}"
export SFT_LORA_PATH="${SFT_LORA_PATH:-${DEFAULT_BASE_DIR}/gen_joint/results_new/scratch_15k}"
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 PYTHONPATH="$PROJECT_ROOT:$OMNIGEN_CODE_ROOT:${PYTHONPATH:-}"

if [[ ! -f "$SFT_LORA_PATH/adapter_config.json" ]]; then
  echo "SFT_LORA_PATH does not contain adapter_config.json: $SFT_LORA_PATH" >&2
  echo "Set SFT_LORA_PATH to the scratch_15k LoRA directory before running eval." >&2
  exit 1
fi

export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-4,5,6,7}

CONFIG_ENTRY=${CONFIG_ENTRY:-"config/grpo.py:general_radiomics_omnigen_4gpu_kl_eval"}
EVAL_LORA_PATH=${EVAL_LORA_PATH:-}

# Generate the eval directory once in the parent shell. If each accelerate
# worker lets eval_omnigen.py create its own timestamp, ranks that start across
# a second boundary can save into different eval_* folders.
EVAL_OUTPUT_ROOT=${EVAL_OUTPUT_ROOT:-"${OUTPUT_DIR:-outputs/test_oldeval/sft_baseline}"}
EVAL_TIMESTAMP=${EVAL_TIMESTAMP:-$(date +%Y.%m.%d_%H.%M.%S)}
EVAL_OUTPUT_DIR=${EVAL_OUTPUT_DIR:-"${EVAL_OUTPUT_ROOT}/eval_${EVAL_TIMESTAMP}"}
mkdir -p "$EVAL_OUTPUT_DIR"
echo "Eval outputs will be written to: $EVAL_OUTPUT_DIR"
echo "SFT LoRA merged into base: $SFT_LORA_PATH"
if [[ -n "$EVAL_LORA_PATH" ]]; then
  if [[ ! -f "$EVAL_LORA_PATH/adapter_config.json" ]]; then
    echo "EVAL_LORA_PATH does not contain adapter_config.json: $EVAL_LORA_PATH" >&2
    exit 1
  fi
  echo "Evaluation LoRA adapter: $EVAL_LORA_PATH"
else
  echo "Evaluation LoRA adapter: <none>"
fi

CMD=(python3 -m accelerate.commands.launch
  --config_file scripts/accelerate_configs/multi_gpu.yaml \
  --num_processes="${NUM_PROCESSES:-4}" \
  --main_process_port "${MAIN_PROCESS_PORT:-29511}" \
  scripts/eval_omnigen.py \
  --config "$CONFIG_ENTRY" \
  --resume_dir "$EVAL_OUTPUT_DIR")

if [[ -n "$EVAL_LORA_PATH" ]]; then
  CMD+=(--eval_lora_path "$EVAL_LORA_PATH")
fi

CMD+=("$@")
"${CMD[@]}"