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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
UMM_ROOT="${UMM_ROOT:-$(cd "${SCRIPT_DIR}/.." && pwd)}"
cd "${UMM_ROOT}"

unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY all_proxy ALL_PROXY

export PYTHONPATH="${UMM_ROOT}/code/umm/runtime:${UMM_ROOT}/code/umm:${UMM_ROOT}/code/umm/stage2_covt:${UMM_ROOT}/third_party/transformers/src:${UMM_ROOT}/third_party/diffusers/src:${PYTHONPATH:-}"
export HF_HOME="${HF_HOME:-${UMM_ROOT}/cache/huggingface}"
export TORCH_HOME="${TORCH_HOME:-${UMM_ROOT}/cache/torch}"
export TMPDIR="${TMPDIR:-${UMM_ROOT}/tmp}"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
export TOKENIZERS_PARALLELISM=false
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"

PYTHON_BIN="${PYTHON_BIN:-${UMM_ROOT}/.env/unvideo/bin/python}"
CHECKPOINT_STEP="${CHECKPOINT_STEP:-2000}"
EXPERIMENT_NAME="${EXPERIMENT_NAME:-stage2_covt_sam_vggt_covt4_v2}"
MODEL_PATH="${MODEL_PATH:-${UMM_ROOT}/models/ckpts/${EXPERIMENT_NAME}/checkpoint-${CHECKPOINT_STEP}-merged}"
MANIFEST="${MANIFEST:-${UMM_ROOT}/data/cvbench_500/validation.jsonl}"
OUTPUT_DIR="${OUTPUT_DIR:-${UMM_ROOT}/outputs/evals/cvbench/${EXPERIMENT_NAME}-checkpoint-${CHECKPOINT_STEP}}"
MASTER_PORT="${MASTER_PORT:-29620}"
SWANLAB_PROJECT="${SWANLAB_PROJECT:-unified-model}"
SWANLAB_RUN_NAME="${SWANLAB_RUN_NAME:-H3_LoRA-ckpt${CHECKPOINT_STEP}-CVBench500}"
PROCESS_NAME="${PROCESS_NAME:-H3_EVAL_${CHECKPOINT_STEP}}"
SWANLAB_KEY_FILE="${SWANLAB_KEY_FILE:-${UMM_ROOT}/secrets/swanlab_key.txt}"

if [[ ! -x "${PYTHON_BIN}" ]]; then
  echo "Missing Python environment: ${PYTHON_BIN}" >&2
  exit 1
fi
if [[ ! -f "${MODEL_PATH}/config.json" ]]; then
  echo "Merged model is not ready: ${MODEL_PATH}" >&2
  exit 1
fi
if [[ ! -s "${MANIFEST}" ]]; then
  echo "CV-Bench manifest is missing: ${MANIFEST}" >&2
  exit 1
fi
if [[ -z "${SWANLAB_API_KEY:-}" ]]; then
  if [[ ! -s "${SWANLAB_KEY_FILE}" ]]; then
    echo "SwanLab key file is missing: ${SWANLAB_KEY_FILE}" >&2
    exit 1
  fi
  IFS= read -r SWANLAB_API_KEY < "${SWANLAB_KEY_FILE}" || true
  if [[ -z "${SWANLAB_API_KEY}" ]]; then
    echo "SwanLab key file is empty: ${SWANLAB_KEY_FILE}" >&2
    exit 1
  fi
  export SWANLAB_API_KEY
fi

CMD=(
  "${PYTHON_BIN}" -m torch.distributed.run
  --standalone
  --nnodes=1
  --nproc-per-node=8
  "--master-port=${MASTER_PORT}"
  "${UMM_ROOT}/code/umm/stage2_covt/eval_cvbench.py"
  --model-path "${MODEL_PATH}"
  --manifest "${MANIFEST}"
  --output-dir "${OUTPUT_DIR}"
  --dtype bfloat16
  --attn-implementation flash_attention_2
  --max-new-tokens 64
  --checkpoint-step "${CHECKPOINT_STEP}"
  --process-name "${PROCESS_NAME}"
  --swanlab
  --swanlab-project "${SWANLAB_PROJECT}"
  --swanlab-run-name "${SWANLAB_RUN_NAME}"
)
CMD+=("$@")

if [[ "${DRY_RUN:-0}" == "1" ]]; then
  printf '%q ' "${CMD[@]}"
  printf '\n'
  exit 0
fi

exec "${CMD[@]}"