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

PROJECT_ROOT=$(CDPATH= cd -- "$(dirname -- "$0")/../.." && pwd)
cd "$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

# CHECKPOINT_ROOT: training checkpoint directory root containing checkpoint-*/lora folders.
CHECKPOINT_ROOT="${CHECKPOINT_ROOT:?Set CHECKPOINT_ROOT to the training checkpoint directory root.}"
CHECKPOINT_GLOB="${CHECKPOINT_GLOB:-checkpoint-*/lora}"
# TEST_JSONL: held-out metadata jsonl used by eval generation and metrics pairing.
TEST_JSONL="${TEST_JSONL:-/home/wenting/zr/wt_dataset/LIDC_IDRI/anno/cxr_synth_anno_mask_test.jsonl}"
# SWEEP_OUTPUT_ROOT: generated images, metrics, per-checkpoint logs, and aggregate summaries.
SWEEP_OUTPUT_ROOT="${SWEEP_OUTPUT_ROOT:-outputs/joint-plan2_1/checkpoint_sweep_$(date +%Y%m%d_%H%M%S)}"
MAX_SAMPLES="${MAX_SAMPLES:-}"
NUM_PROCESSES="${NUM_PROCESSES:-4}"
FORCE="${FORCE:-0}"
ALLOW_METRIC_SKIP="${ALLOW_METRIC_SKIP:-0}"
ALLOW_MISSING_OUTPUTS="${ALLOW_MISSING_OUTPUTS:-0}"
RESIZE_POLICY="${RESIZE_POLICY:-none}"
METRICS_BATCH_SIZE="${METRICS_BATCH_SIZE:-32}"
mkdir -p "$SWEEP_OUTPUT_ROOT/logs"
SUMMARY_CSV="$SWEEP_OUTPUT_ROOT/sweep_summary.csv"
SUMMARY_JSON="$SWEEP_OUTPUT_ROOT/sweep_summary.json"
FAILED="$SWEEP_OUTPUT_ROOT/failed_checkpoints.txt"
echo 'checkpoint,output_dir,metrics_report,psnr_mean,ssim_mean,lpips_mean,fid_value,num_scored,num_missing_or_failed' > "$SUMMARY_CSV"
echo '[]' > "$SUMMARY_JSON"
: > "$FAILED"

mapfile -t CHECKPOINTS < <(find "$CHECKPOINT_ROOT" -path "*/${CHECKPOINT_GLOB}" -type d | sort)
for CKPT in "${CHECKPOINTS[@]}"; do
  STEP_NAME=$(basename "$(dirname "$CKPT")")
  OUT_DIR="$SWEEP_OUTPUT_ROOT/$STEP_NAME/generated"
  METRICS_DIR="$SWEEP_OUTPUT_ROOT/$STEP_NAME/metrics"
  LOG_FILE="$SWEEP_OUTPUT_ROOT/logs/${STEP_NAME}.log"
  if [ -f "$METRICS_DIR/metrics_report.json" ] && [ "$FORCE" != "1" ]; then
    echo "Skipping existing $STEP_NAME"
  else
    mkdir -p "$OUT_DIR" "$METRICS_DIR"
    echo "Evaluating $CKPT" | tee "$LOG_FILE"
    if ! EVAL_LORA_PATH="$CKPT" OUTPUT_DIR="$OUT_DIR" TEST_JSONL="$TEST_JSONL" NUM_PROCESSES="$NUM_PROCESSES" MAX_SAMPLES="$MAX_SAMPLES" EVAL_EXACT_OUTPUT_DIR=1 \
      bash scripts/single_node/eval_4gpu_joint_plan2_1.sh >> "$LOG_FILE" 2>&1; then
      echo "$CKPT" >> "$FAILED"
      continue
    fi
    if ! GENERATED_DIR="$OUT_DIR" METRICS_OUT="$METRICS_DIR" TEST_JSONL="$TEST_JSONL" MAX_SAMPLES="$MAX_SAMPLES" ALLOW_METRIC_SKIP="$ALLOW_METRIC_SKIP" ALLOW_MISSING_OUTPUTS="$ALLOW_MISSING_OUTPUTS" RESIZE_POLICY="$RESIZE_POLICY" METRICS_BATCH_SIZE="$METRICS_BATCH_SIZE" \
      bash scripts/single_node/metrics_for_eval_outputs_joint_plan2_1.sh >> "$LOG_FILE" 2>&1; then
      echo "$CKPT metrics" >> "$FAILED"
      continue
    fi
  fi
  if [ -f "$METRICS_DIR/metrics_report.json" ]; then
    if ! python3 - "$CKPT" "$OUT_DIR" "$METRICS_DIR/metrics_report.json" "$SUMMARY_CSV" "$SUMMARY_JSON" <<'PY'
import csv, json, sys
ckpt, out_dir, report_path, csv_path, json_path = sys.argv[1:]
report = json.load(open(report_path))
if not report.get('allow_metric_skip', False):
    if report.get('lpips_mean') is None or report.get('fid_value') is None:
        raise SystemExit(f"Official metrics incomplete for {ckpt}: lpips_mean={report.get('lpips_mean')}, fid_value={report.get('fid_value')}")
if not report.get('allow_missing_outputs', False) and report.get('num_missing_or_failed', 0):
    raise SystemExit(f"Missing outputs for {ckpt}: {report.get('num_missing_or_failed')}")
row = {
    'checkpoint': ckpt,
    'output_dir': out_dir,
    'metrics_report': report_path,
    'psnr_mean': report.get('psnr_mean'),
    'ssim_mean': report.get('ssim_mean'),
    'lpips_mean': report.get('lpips_mean'),
    'fid_value': report.get('fid_value'),
    'num_scored': report.get('num_scored'),
    'num_missing_or_failed': report.get('num_missing_or_failed'),
}
with open(csv_path, 'a', newline='') as f:
    writer = csv.DictWriter(f, fieldnames=list(row))
    writer.writerow(row)
items = json.load(open(json_path))
items.append(row)
json.dump(items, open(json_path, 'w'), indent=2)
PY
    then
      echo "$CKPT metrics_summary" >> "$FAILED"
      continue
    fi
  fi
done