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

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

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_BCD}"
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:-29671}"

CONFIG="${UMM_ROOT}/configs/stage2_covt_from3k_bcd_2k.yaml"
CONFIG_TEMPLATE="${UMM_ROOT}/configs/stage2_covt_from3k_bcd_2k.template.yaml"
SOURCE_CKPT="${UMM_ROOT}/outputs/stage2_covt_sam_vggt_covt4_v2/checkpoint-3000"
OUTPUT_DIR="${UMM_ROOT}/models/ckpts/stage2_covt_from3k_bcd_2k"
BOOTSTRAP_LINK="${OUTPUT_DIR}/checkpoint-3000"
LOG_FILE="${UMM_ROOT}/logs/H3_LoRA_from3k_BCD_train.log"

if [[ "${CUDA_VISIBLE_DEVICES}" != "0,1,2,3,4,5,6,7" ]]; then
  echo "Formal B/C/D 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 unavailable." >&2
  exit 1
fi
if [[ ! -f "${CONFIG}" ]]; then
  "${UMM_ROOT}/.env/unvideo/bin/python" "${UMM_ROOT}/scripts/render_config.py" \
    --template "${CONFIG_TEMPLATE}" \
    --output "${CONFIG}"
fi
if [[ ! -s "${SWANLAB_KEY_FILE}" ]]; then
  echo "SwanLab key file is missing or empty: ${SWANLAB_KEY_FILE}" >&2
  exit 1
fi
for required in trainer_state.json adapter_model.safetensors non_lora_state_dict.bin latest; do
  if [[ ! -e "${SOURCE_CKPT}/${required}" ]]; then
    echo "Source checkpoint is incomplete: ${SOURCE_CKPT}/${required}" >&2
    exit 1
  fi
done

mkdir -p "${OUTPUT_DIR}" "${UMM_ROOT}/logs"
if ! find "${OUTPUT_DIR}" -maxdepth 1 -type d -name 'checkpoint-*' -exec test -f '{}/trainer_state.json' ';' -print -quit | grep -q .; then
  if [[ -L "${BOOTSTRAP_LINK}" ]]; then
    if [[ "$(readlink -f "${BOOTSTRAP_LINK}")" != "$(readlink -f "${SOURCE_CKPT}")" ]]; then
      echo "Unexpected checkpoint-3000 symlink target: ${BOOTSTRAP_LINK}" >&2
      exit 1
    fi
  elif [[ -e "${BOOTSTRAP_LINK}" ]]; then
    echo "Refusing to overwrite existing non-symlink path: ${BOOTSTRAP_LINK}" >&2
    exit 1
  else
    ln -s "${SOURCE_CKPT}" "${BOOTSTRAP_LINK}"
  fi
fi

"${UMM_ROOT}/.env/unvideo/bin/python" - "${CONFIG}" "${SOURCE_CKPT}" <<'PY'
import json
import sys
from pathlib import Path

import yaml

config_path, checkpoint_path = map(Path, sys.argv[1:])
config = yaml.safe_load(config_path.read_text(encoding="utf-8"))
state = json.loads((checkpoint_path / "trainer_state.json").read_text(encoding="utf-8"))
assert int(state["global_step"]) == 3000, state["global_step"]
assert [config["data"][key] for key in ("stage_0_step", "stage_1_step", "stage_2_step")] == [3000, 5000, 7000]
assert config["training"]["max_steps"] == 9000
assert config["training"]["cvbench_eval_every_n_steps"] == 500
assert config["training"]["save_steps"] == 1000
assert config["training"]["save_total_limit"] >= 6
optim = sorted((checkpoint_path / "global_step3000").glob("*_optim_states.pt"))
assert len(optim) == 8, len(optim)
print("[H3_LoRA_BCD] preflight: source=3000, B=[3000,5000), C=[5000,7000), D=[7000,9000), val=500")
PY

CONFIG="${CONFIG}" DRY_RUN=1 bash code/umm/stage2_covt/run_train.sh \
  > "${UMM_ROOT}/logs/H3_LoRA_from3k_BCD_dry_run.log"

echo "[H3_LoRA_BCD] starting resumable 8-GPU formal training" | tee -a "${LOG_FILE}"
CONFIG="${CONFIG}" bash code/umm/stage2_covt/run_train.sh \
  2>&1 | tee -a "${LOG_FILE}"