#!/usr/bin/env bash
# Bundle weights + examples for Hugging Face Space upload (standalone, no vendor/).
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VYOLO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
ULTRA_SRC="${VYOLO_ROOT}/ultralytics"
SPACE_DIR="${SCRIPT_DIR}"
WEIGHTS_SRC="${ULTRA_SRC}/Proposed/yolo11_training/HRIPCB_Results/yolo11n_hripcb_200epochs_batch16/weights/best.pt"
SAMPLES_SRC="${ULTRA_SRC}/Proposed/yolo11_training/HRIPCB_Results/yolo11n_hripcb_200epochs_batch16/prebackbone_samples"

echo "==> Preparing Hugging Face Space in ${SPACE_DIR}"

# 1) Weights — extract prebackbone-only (~few MB) from best.pt
mkdir -p "${SPACE_DIR}/weights"
PB_ONLY="${SPACE_DIR}/weights/prebackbone_a11_ca.pt"
if [[ -f "${WEIGHTS_SRC}" ]]; then
  PYTHON="${PYTHON:-python3}"
  if "${PYTHON}" -c "import torch" 2>/dev/null; then
    export ULTRALYTICS_ROOT="${ULTRA_SRC}"
    (cd "${SPACE_DIR}" && "${PYTHON}" extract_prebackbone_weights.py --ckpt "${WEIGHTS_SRC}" --out "${PB_ONLY}")
    echo "    Prebackbone-only weights: $(du -h "${PB_ONLY}" | cut -f1)"
    if [[ "${INCLUDE_FULL_CKPT:-0}" == "1" ]]; then
      cp -f "${WEIGHTS_SRC}" "${SPACE_DIR}/weights/best.pt"
      echo "    Also copied full best.pt (optional fallback)"
    fi
  else
    echo "    WARN: torch not available — place prebackbone_a11_ca.pt in weights/ manually"
  fi
else
  echo "    WARN: ${WEIGHTS_SRC} not found — place prebackbone_a11_ca.pt in weights/"
fi

# 2) Example pairs — all *_input.jpg / *_reference.jpg under prebackbone_samples
EX_DIR="${SPACE_DIR}/examples"
rm -rf "${EX_DIR}"
mkdir -p "${EX_DIR}"
_n_pairs=0
if [[ -d "${SAMPLES_SRC}" ]]; then
  while IFS= read -r -d '' inp; do
    ref="${inp/_input./_reference.}"
    if [[ ! -f "${ref}" ]]; then
      continue
    fi
    base="$(basename "${inp}")"
    cp -f "${inp}" "${EX_DIR}/${base}"
    cp -f "${ref}" "${EX_DIR}/${base/_input./_reference.}"
    _n_pairs=$((_n_pairs + 1))
  done < <(find "${SAMPLES_SRC}" -name '*_input.jpg' -print0 | sort -z)
fi
if [[ "${_n_pairs}" -gt 0 ]]; then
  echo "    Examples: ${_n_pairs} pairs ($(du -sh "${EX_DIR}" | cut -f1))"
else
  echo "    WARN: no prebackbone_samples found for examples"
fi

# 3) Drop legacy vendored ultralytics if present
if [[ -d "${SPACE_DIR}/vendor" ]]; then
  rm -rf "${SPACE_DIR}/vendor"
  echo "    Removed legacy vendor/ (standalone demo)"
fi

TOTAL="$(du -sh "${SPACE_DIR}" | cut -f1)"
echo "==> Done (total ${TOTAL}). Next:"
echo "    cd ${SPACE_DIR}"
echo "    git init && git lfs install && git lfs track '*.pt'"
echo "    git add . && git commit -m 'RefDiffNet standalone demo'"
echo "    git remote add origin https://huggingface.co/spaces/USER/SPACE"
echo "    git push"