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

PROJECT_ROOT=$(CDPATH= cd -- "$(dirname -- "$0")/../.." && pwd)
cd "$PROJECT_ROOT"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)

# Optional W&B credentials. This file may export WANDB_API_KEY; the key is not printed.
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/zr/gen_code_plan2_1}"
export CONFIG_ENTRY="${CONFIG_ENTRY:-config/grpo.py:general_radiomics_joint_omnigen_4gpu_cm_kl_eval}"
# SFT_LORA_PATH: SFT adapter folder merged before eval RL LoRA is loaded.
export SFT_LORA_PATH="${SFT_LORA_PATH:-/home/wenting/zr/gen_code_plan2_1/results_new/scratch_15k}"
# MASK_MODULES_PATH: companion mask_modules.bin for the joint model.
export MASK_MODULES_PATH="${MASK_MODULES_PATH:-/home/wenting/zr/gen_code_plan2_1/results_new/scratch_15k/mask_modules.bin}"
# TRAIN_JSONL: training metadata jsonl; printed for provenance, not consumed by eval generation.
export TRAIN_JSONL="${TRAIN_JSONL:-/home/wenting/zr/wt_dataset/LIDC_IDRI/anno/cxr_synth_anno_mask_train.jsonl}"
# TEST_JSONL: held-out metadata jsonl used by eval generation.
export TEST_JSONL="${TEST_JSONL:-/home/wenting/zr/wt_dataset/LIDC_IDRI/anno/cxr_synth_anno_mask_test.jsonl}"
# SAVE_DIR: training run directory that usually contains checkpoints.
export SAVE_DIR="${SAVE_DIR:-logs/radiomics/joint-scratch15k-lung-beta0p005}"
# OUTPUT_DIR: root for eval-generated images and downstream metrics.
export OUTPUT_DIR="${OUTPUT_DIR:-outputs/joint-plan2_1/eval_noise001_${TIMESTAMP}}"
export EVAL_LORA_PATH="${EVAL_LORA_PATH:-logs/radiomics/joint-scratch15k-lung-beta0p005/checkpoints/checkpoint-60/lora}"
export EVAL_EXACT_OUTPUT_DIR="${EVAL_EXACT_OUTPUT_DIR:-0}" # If set, eval generation will write to OUTPUT_DIR without timestamp subfolder and will look for LoRA in OUTPUT_DIR; otherwise, OUTPUT_DIR is a parent folder for timestamped eval outputs and EVAL_LORA_PATH must be set explicitly.——如果是1就是断点续跑

# Reward controls are printed for provenance; eval generation does not compute reward.
export REWARD_MASK_CHANNELS="${REWARD_MASK_CHANNELS:-0,1}"
export REWARD_FEATURE_PREFIXES="${REWARD_FEATURE_PREFIXES:-cm_}"
export REWARD_BASE_FEATURE_FAMILIES="${REWARD_BASE_FEATURE_FAMILIES:-cm}"
export REWARD_DISTANCE_MODE="${REWARD_DISTANCE_MODE:-relative_l1}"
export REWARD_SCALE="${REWARD_SCALE:-1.0}"
export REWARD_MISSING_SCORE="${REWARD_MISSING_SCORE:--10.0}"

# Sampling controls. NOISE_LEVEL should remain 0.0 for deterministic eval unless intentionally overridden.
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 SAME_LATENT="${SAME_LATENT:-false}"
export GLOBAL_STD="${GLOBAL_STD:-true}"
export MAX_EVAL_BATCHES="${MAX_EVAL_BATCHES:-2}" # cy

# Training/GRPO controls are part of config provenance; eval mainly uses LoRA/checkpoint paths and sampling controls.
export NUM_IMAGE_PER_PROMPT="${NUM_IMAGE_PER_PROMPT:-4}"
export KL_BETA="${KL_BETA:-0.005}"
export LEARNING_RATE="${LEARNING_RATE:-1e-4}"
export TRAIN_BATCH_SIZE="${TRAIN_BATCH_SIZE:-}"
export GRAD_ACCUM="${GRAD_ACCUM:-}"
export MAX_TRAIN_STEPS="${MAX_TRAIN_STEPS:-}"
export SAVE_FREQ="${SAVE_FREQ:-5}"
export EVAL_FREQ="${EVAL_FREQ:-5}"
export RL_LORA_RANK="${RL_LORA_RANK:-32}"
export RL_LORA_ALPHA="${RL_LORA_ALPHA:-64}"
export RL_LORA_TARGET_MODULES="${RL_LORA_TARGET_MODULES:-qkv_proj,o_proj}"

export NUM_PROCESSES="${NUM_PROCESSES:-4}"
export MAIN_PROCESS_PORT="${MAIN_PROCESS_PORT:-29501}"
export MIXED_PRECISION="${MIXED_PRECISION:-bf16}"
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}"
export WANDB_PROJECT="${WANDB_PROJECT:-flow_grpo}"
export WANDB_NAME="${WANDB_NAME:-joint-plan2_1-eval-${TIMESTAMP}}"
export WANDB_MODE="${WANDB_MODE:-disabled}"
export PYTHONPATH="$PROJECT_ROOT:$OMNIGEN_CODE_ROOT:${PYTHONPATH:-}"

mkdir -p "$OUTPUT_DIR"
export LAUNCH_CONFIG="$OUTPUT_DIR/launch_config.txt"
{
  echo "timestamp=${TIMESTAMP}"
  echo "repo_path=${PROJECT_ROOT}"
  echo "config_entry=${CONFIG_ENTRY}"
  echo "omnigen_code_root=${OMNIGEN_CODE_ROOT}"
  echo "sft_lora_path=${SFT_LORA_PATH}"
  echo "mask_modules_path=${MASK_MODULES_PATH}"
  echo "train_jsonl=${TRAIN_JSONL}"
  echo "test_jsonl=${TEST_JSONL}"
  echo "save_dir=${SAVE_DIR}"
  echo "output_dir=${OUTPUT_DIR}"
  echo "eval_lora_path=${EVAL_LORA_PATH}"
  echo "eval_exact_output_dir=${EVAL_EXACT_OUTPUT_DIR}"
  echo "max_samples=${MAX_SAMPLES:-}"
  echo "launch_config=${LAUNCH_CONFIG}"
  echo "reward_mask_channels=${REWARD_MASK_CHANNELS}"
  echo "reward_feature_prefixes=${REWARD_FEATURE_PREFIXES}"
  echo "reward_base_feature_families=${REWARD_BASE_FEATURE_FAMILIES}"
  echo "reward_distance_mode=${REWARD_DISTANCE_MODE}"
  echo "reward_scale=${REWARD_SCALE}"
  echo "reward_missing_score=${REWARD_MISSING_SCORE}"
  echo "num_steps=${NUM_STEPS}"
  echo "eval_num_steps=${EVAL_NUM_STEPS}"
  echo "noise_level=${NOISE_LEVEL}"
  echo "sde_type=${SDE_TYPE}"
  echo "guidance_scale=${GUIDANCE_SCALE}"
  echo "img_guidance_scale=${IMG_GUIDANCE_SCALE}"
  echo "same_latent=${SAME_LATENT}"
  echo "global_std=${GLOBAL_STD}"
  echo "max_eval_batches=${MAX_EVAL_BATCHES}"
  echo "num_image_per_prompt=${NUM_IMAGE_PER_PROMPT}"
  echo "kl_beta=${KL_BETA}"
  echo "learning_rate=${LEARNING_RATE}"
  echo "train_batch_size=${TRAIN_BATCH_SIZE}"
  echo "grad_accum=${GRAD_ACCUM}"
  echo "max_train_steps=${MAX_TRAIN_STEPS}"
  echo "save_freq=${SAVE_FREQ}"
  echo "eval_freq=${EVAL_FREQ}"
  echo "rl_lora_rank=${RL_LORA_RANK}"
  echo "rl_lora_alpha=${RL_LORA_ALPHA}"
  echo "rl_lora_target_modules=${RL_LORA_TARGET_MODULES}"
  echo "num_processes=${NUM_PROCESSES}"
  echo "main_process_port=${MAIN_PROCESS_PORT}"
  echo "mixed_precision=${MIXED_PRECISION}"
  echo "cuda_visible_devices=${CUDA_VISIBLE_DEVICES}"
  echo "wandb_project=${WANDB_PROJECT}"
  echo "wandb_name=${WANDB_NAME}"
  echo "wandb_mode=${WANDB_MODE}"
  if [[ -n "${WANDB_API_KEY:-}" ]]; then echo "wandb_api_key_set=true"; else echo "wandb_api_key_set=false"; fi
  echo "wandb_env_file=${WANDB_ENV_FILE}"
} | 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"
  --output_dir "$OUTPUT_DIR"
)
if [ "$EVAL_EXACT_OUTPUT_DIR" = "1" ]; then
  CMD+=(--resume_dir "$OUTPUT_DIR")
fi
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+=("$@")
"${CMD[@]}"