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

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

# Downloads and SwanLab traffic must go directly from this host.
unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY all_proxy ALL_PROXY

export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
export UMM_PROCESS_NAME="${UMM_PROCESS_NAME:-H3_LoRA}"
export SWANLAB_PROJ_NAME="${SWANLAB_PROJ_NAME:-unified-model}"
export SWANLAB_KEY_FILE="${SWANLAB_KEY_FILE:-${UMM_ROOT}/secrets/swanlab_key.txt}"
export MASTER_PORT="${MASTER_PORT:-29617}"

if [[ "${CUDA_VISIBLE_DEVICES}" != "0,1,2,3,4,5,6,7" ]]; then
  echo "H3_LoRA formal launch requires all eight GPUs; got ${CUDA_VISIBLE_DEVICES}" >&2
  exit 1
fi
if [[ ! -x "${UMM_ROOT}/.env/unvideo/bin/python" ]]; then
  echo "Portable environment is not deployed." >&2
  exit 1
fi
if [[ ! -s "${SWANLAB_KEY_FILE}" ]]; then
  echo "SwanLab key file is missing or empty: ${SWANLAB_KEY_FILE}" >&2
  exit 1
fi
if [[ ! -s "${UMM_ROOT}/data/cvbench_500/validation.jsonl" ]]; then
  echo "CV-Bench validation subset is not ready." >&2
  exit 1
fi

mkdir -p outputs logs
"${UMM_ROOT}/.env/unvideo/bin/python" "${UMM_ROOT}/scripts/validate_covt_references.py" \
  2>&1 | tee "${UMM_ROOT}/logs/covt_reference_closure.log"
"${UMM_ROOT}/.env/unvideo/bin/python" "${UMM_ROOT}/scripts/verify_bundle.py" --materialized \
  2>&1 | tee "${UMM_ROOT}/logs/H3_LoRA_preflight.log"
nvidia-smi \
  --query-gpu=index,name,memory.total,memory.used,memory.free,utilization.gpu,temperature.gpu \
  --format=csv,noheader \
  | tee logs/H3_LoRA_gpu_before_launch.log
sha256sum \
  configs/stage2_covt.yaml \
  data/cvbench_500/validation.jsonl \
  data/cvbench_500/summary.json \
  code/umm/stage2_covt/train.py \
  code/umm/stage2_covt/cvbench_eval.py \
  > logs/H3_LoRA_launch_fingerprint.sha256
DRY_RUN=1 bash code/umm/stage2_covt/run_train.sh \
  --run_name H3_LoRA-dry-run \
  > logs/H3_LoRA_dry_run.log

SMOKE_DIR="${UMM_ROOT}/outputs/H3_LoRA_8gpu_smoke"
SMOKE_SENTINEL="${SMOKE_DIR}/SMOKE_COMPLETE"
if [[ "${FORCE_SMOKE:-0}" == "1" || ! -f "${SMOKE_SENTINEL}" ]]; then
  mkdir -p "${SMOKE_DIR}"
  echo "[H3_LoRA] starting 8-GPU one-step smoke with full CV-Bench-500 evaluation"
  bash code/umm/stage2_covt/run_train.sh \
    --output_dir "${SMOKE_DIR}" \
    --max_steps 1 \
    --per_device_train_batch_size 1 \
    --save_strategy no \
    --run_name H3_LoRA-8gpu-smoke \
    --cvbench_eval_every_n_steps 1 \
    2>&1 | tee "${UMM_ROOT}/logs/H3_LoRA_8gpu_smoke.log"
  date -u +%FT%TZ > "${SMOKE_SENTINEL}"
fi

CAPACITY_DIR="${UMM_ROOT}/outputs/H3_LoRA_8gpu_batch4_smoke"
CAPACITY_SENTINEL="${CAPACITY_DIR}/SMOKE_COMPLETE"
if [[ "${FORCE_SMOKE:-0}" == "1" || ! -f "${CAPACITY_SENTINEL}" ]]; then
  mkdir -p "${CAPACITY_DIR}"
  echo "[H3_LoRA] starting 8-GPU batch-4 capacity smoke"
  bash code/umm/stage2_covt/run_train.sh \
    --output_dir "${CAPACITY_DIR}" \
    --max_steps 1 \
    --per_device_train_batch_size 4 \
    --save_strategy no \
    --run_name H3_LoRA-8gpu-batch4-smoke \
    --cvbench_eval_every_n_steps 0 \
    2>&1 | tee "${UMM_ROOT}/logs/H3_LoRA_8gpu_batch4_smoke.log"
  date -u +%FT%TZ > "${CAPACITY_SENTINEL}"
fi

echo "[H3_LoRA] connectivity/eval and batch-4 capacity smokes passed; starting resumable 8-GPU formal training"
bash code/umm/stage2_covt/run_train.sh \
  --run_name H3_LoRA-stage2-covt4-8gpu \
  2>&1 | tee -a "${UMM_ROOT}/logs/H3_LoRA_train.log"