#!/usr/bin/env bash
# 从 Label Studio 的最新标注重新训练检测模型，并让 ML backend 用上新权重。
#
#   bash scripts/retrain_from_label_studio.sh
#
# 也由 ML backend 在点击 Start Training 时自动调用（输出见 $RETRAIN_LOG）。
# 四步：导出标注 -> 装配数据集 -> 训练 -> 切换权重并重启后端。
set -euo pipefail

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

PROJECT="${LS_PROJECT_ID:-39}"
IMAGES="${LS_IMAGES_DIR:-$REPO/runs/boxes_shap_label_studio/images}"
STAMP=$(date +%Y%m%d_%H%M)
EXPORT="${LS_EXPORT_DIR:-/tmp}/ls_export_${STAMP}.zip"
DATASET="$REPO/data/film_yolo_${STAMP}"
WORK_DIR="$REPO/runs/detect/film_${STAMP}"

PYTHON="${REPO_PYTHON:-$REPO/.venv/bin/python}"
# 用跑着 Label Studio 的那个环境导出，确保读的是同一个库
LS_BIN="${LS_BIN:-/home/deep/.v/label/bin/label-studio}"
[ -x "$LS_BIN" ] || LS_BIN="$REPO/.venv/bin/label-studio"
export LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT="${LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT:-/home/deep/labeling}"

# 训练参数的额外覆盖，冒烟测试时用，如 RETRAIN_EXTRA_SET="train.epochs=1 train.fraction=0.02"
EXTRA_SET="${RETRAIN_EXTRA_SET:-}"

# 训完要不要重启 ML backend。被后端调用时必须重启（权重只在启动时加载一次）；
# 纯手动路线不跑后端的话设 0，免得凭空起一个常驻进程占显存
RESTART_BACKEND="${RESTART_BACKEND:-1}"

echo "===== [$(date '+%F %T')] 开始重训（项目 $PROJECT）====="

echo "--- 1/4 导出标注 -> $EXPORT"
"$LS_BIN" export "$PROJECT" YOLO --export-path="$EXPORT"

echo "--- 2/4 装配数据集 -> $DATASET"
"$PYTHON" scripts/build_yolo_dataset.py \
    --export "$EXPORT" --images "$IMAGES" --out "$DATASET" --overwrite

echo "--- 3/4 训练 -> $WORK_DIR"
# shellcheck disable=SC2086  # EXTRA_SET 需要按空格拆成多个参数
"$PYTHON" detect.py --config configs/detect/film.yaml \
    --work-dir "$WORK_DIR" \
    --set data.yaml="$DATASET/data.yaml" $EXTRA_SET

echo "--- 4/4 切换权重"
# 指向最新权重的稳定软链接：启动脚本与 push_predictions 都读它，换模型不用改任何脚本
ln -sfn "$WORK_DIR/weights/best.pt" "$REPO/runs/detect/current.pt"

if [ "$RESTART_BACKEND" = "1" ]; then
    echo "--- 重启 ML backend"
    # 必须显式传 YOLO_WEIGHTS 并清掉继承来的变量：本脚本可能是旧后端的子进程，
    # 它环境里已经 export 了 YOLO_WEIGHTS（旧权重）和可能的 RETRAIN_EXTRA_SET（冒烟参数），
    # 不清掉的话新后端会继续用旧模型、下次训练也会沿用冒烟参数
    env -u RETRAIN_EXTRA_SET YOLO_WEIGHTS="$REPO/runs/detect/current.pt" \
        bash "$REPO/scripts/start_ml_backend.sh"
else
    echo "--- 跳过重启 ML backend（RESTART_BACKEND=0）"
fi

echo "===== [$(date '+%F %T')] 重训完成：$WORK_DIR ====="
echo "指标见 $WORK_DIR/metrics.json，曲线见 $WORK_DIR/results.png"