#!/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
TIMESTAMP=$(date +%Y%m%d_%H%M%S)

export OMNIGEN_CODE_ROOT="${OMNIGEN_CODE_ROOT:-/apdcephfs_zwfy2/share_303944931/connorxian/Code/medUni/wt/zr/gen_joint}"
export CONFIG_ENTRY="${CONFIG_ENTRY:-config/grpo.py:general_radiomics_omnigen_4gpu_kl_eval}"
export SFT_LORA_PATH="${SFT_LORA_PATH:-/apdcephfs_zwfy2/share_303944931/connorxian/Code/medUni/wt/zr/gen_joint/results_new/scratch_15k}"
export DATASET_ROOT="${DATASET_ROOT:-/apdcephfs_zwfy2/share_303944931/connorxian/Code/medUni/wt/zr/dataset}"
export TRAIN_JSONL="${TRAIN_JSONL:-/apdcephfs_zwfy2/share_303944931/connorxian/Code/medUni/wt/zr/dataset/anno/cxr_synth_anno_mask_train.jsonl}"
export TEST_JSONL="${TEST_JSONL:-/apdcephfs_zwfy2/share_303944931/connorxian/Code/medUni/wt/zr/dataset/anno/cxr_synth_anno_mask_test.jsonl}"
export DATASET_PATH_REMAP_FROM="${DATASET_PATH_REMAP_FROM:-/home/wenting/dataset/LIDC_IDRI}"
export DATASET_PATH_REMAP_TO="${DATASET_PATH_REMAP_TO:-$DATASET_ROOT}"
export MIRP_REPO_ROOT="${MIRP_REPO_ROOT:-/apdcephfs_zwfy2/share_303944931/connorxian/Code/medUni/wt/zr/mirp}"
export EVAL_LORA_PATH="${EVAL_LORA_PATH:-}"
export OUTPUT_DIR="${OUTPUT_DIR:-outputs/img_only_scratch30k/eval_${TIMESTAMP}}"

export NUM_STEPS="${NUM_STEPS:-10}"
export EVAL_NUM_STEPS="${EVAL_NUM_STEPS:-50}"
export NOISE_LEVEL="${NOISE_LEVEL:-0.0}"
export SDE_TYPE="${SDE_TYPE:-cps}"
export GUIDANCE_SCALE="${GUIDANCE_SCALE:-2.5}"
export IMG_GUIDANCE_SCALE="${IMG_GUIDANCE_SCALE:-2.0}"
export TRAIN_BATCH_SIZE="${TRAIN_BATCH_SIZE:-32}"
export TEST_BATCH_SIZE="${TEST_BATCH_SIZE:-32}"
export NUM_PROCESSES="${NUM_PROCESSES:-4}"
export MAIN_PROCESS_PORT="${MAIN_PROCESS_PORT:-29501}"
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
export HF_HOME="${HF_HOME:-/apdcephfs_zwfy2/share_303944931/connorxian/Code/medUni/wt/zr/.cache/huggingface}"
export HF_HUB_CACHE="${HF_HUB_CACHE:-/apdcephfs_zwfy2/share_303944931/connorxian/Code/medUni/wt/zr/.cache/huggingface/hub}"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}"
export WANDB_MODE="${WANDB_MODE:-disabled}"
export WANDB_PROJECT="${WANDB_PROJECT:-flow_grpo_img_only_scratch30k}"
export WANDB_GROUP="${WANDB_GROUP:-manual_img_only_eval}"
export WANDB_NAME="${WANDB_NAME:-${EXP_NAME:-img-only-scratch30k-eval}-${TIMESTAMP}}"
export WANDB_TAGS="${WANDB_TAGS:-img-only,scratch30k,eval}"
export REWARD_MASK_CHANNELS="${REWARD_MASK_CHANNELS:-0,1}"
export PYTHONPATH="$PROJECT_ROOT:$OMNIGEN_CODE_ROOT:${PYTHONPATH:-}"

mkdir -p "$OUTPUT_DIR" "$HF_HUB_CACHE"
LAUNCH_CONFIG="$OUTPUT_DIR/launch_config.txt"
{
  echo "timestamp=${TIMESTAMP}"
  echo "repo_path=${PROJECT_ROOT}"
  echo "config_entry=${CONFIG_ENTRY}"
  echo "server_profile=${SERVER_PROFILE:-}"
  echo "omnigen_code_root=${OMNIGEN_CODE_ROOT}"
  echo "sft_lora_path=${SFT_LORA_PATH}"
  echo "dataset_root=${DATASET_ROOT}"
  echo "train_jsonl=${TRAIN_JSONL}"
  echo "test_jsonl=${TEST_JSONL}"
  echo "dataset_path_remap_from=${DATASET_PATH_REMAP_FROM}"
  echo "dataset_path_remap_to=${DATASET_PATH_REMAP_TO}"
  echo "mirp_repo_root=${MIRP_REPO_ROOT}"
  echo "eval_lora_path=${EVAL_LORA_PATH}"
  echo "output_dir=${OUTPUT_DIR}"
  echo "eval_num_steps=${EVAL_NUM_STEPS}"
  echo "noise_level=${NOISE_LEVEL}"
  echo "sde_type=${SDE_TYPE}"
  echo "reward_mask_channels=${REWARD_MASK_CHANNELS}"
  echo "num_processes=${NUM_PROCESSES}"
  echo "wandb_project=${WANDB_PROJECT}"
  echo "wandb_group=${WANDB_GROUP}"
  echo "wandb_name=${WANDB_NAME}"
  echo "wandb_tags=${WANDB_TAGS}"
  echo "wandb_mode=${WANDB_MODE}"
} | tee "$LAUNCH_CONFIG"

CMD=(python3 -m accelerate.commands.launch
  --config_file scripts/accelerate_configs/multi_gpu.yaml
  --num_processes="$NUM_PROCESSES"
  --main_process_port "$MAIN_PROCESS_PORT"
  scripts/eval_omnigen.py
  --config "$CONFIG_ENTRY"
  --resume_dir "$OUTPUT_DIR")

if [[ -n "$EVAL_LORA_PATH" ]]; then
  CMD+=(--eval_lora_path "$EVAL_LORA_PATH")
fi
if [[ -n "${MAX_SAMPLES:-}" ]]; then
  CMD+=(--max_samples "$MAX_SAMPLES")
fi

"${CMD[@]}"