#!/usr/bin/env bash
# 从 Label Studio 的最新标注重新训练检测模型，并让 ML backend 用上新权重。
#
#   bash scripts/retrain_from_label_studio.sh
#
# 也由 ML backend 在点击 Start Training 时自动调用（输出见 $RETRAIN_LOG）。
# 五步：导出标注 -> 装配数据集 -> 切成固定边长的片 -> 训练 -> 切换权重、归档、重启后端。
#
# 产线要求模型输入固定为 TILE_SIZE×TILE_SIZE 且禁止 resize，所以训练数据必须先切片；
# 推理侧（可视化、批量推送、ML backend）也要用同一个尺寸，否则输入分布对不上。
# TILE_SIZE=0 可跳过切片，训练未切片的数据集（只用于和旧模型对比）。
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)

# 切片参数。改尺寸时这里与 configs/detect/film.yaml 的 imgsz 会一起生效（见第 4 步）
TILE_SIZE="${TILE_SIZE:-128}"
TILE_OVERLAP="${TILE_OVERLAP:-32}"
# >0.5 保证一个框最多出现在一片里，不会在相邻片各留一个残框
MIN_VISIBILITY="${MIN_VISIBILITY:-0.6}"

# 边长写进目录名：同时留着 128 和 160 两套时，一眼能分清谁是谁
SIZE_TAG=$([ "$TILE_SIZE" -gt 0 ] && echo "$TILE_SIZE" || echo "raw")
EXPORT="${LS_EXPORT_DIR:-/tmp}/ls_export_${STAMP}.zip"
DATASET="$REPO/data/film_yolo_${STAMP}"
TILED="$REPO/data/film_yolo_${SIZE_TAG}_${STAMP}"
WORK_DIR="$REPO/runs/detect/film_${SIZE_TAG}_${STAMP}"
# 归档目录：每次训练留一份 film_<边长>_<时间戳>.pt 与同名的 OpenVINO 导出，
# 供部署机取用。软链接会被下次训练改写，这里的带戳副本不会。
MODELS_DIR="$REPO/runs/detect/models"
EXPORT_OPENVINO="${EXPORT_OPENVINO:-1}"

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/5 导出标注 -> $EXPORT"
"$LS_BIN" export "$PROJECT" YOLO --export-path="$EXPORT"

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

if [ "$TILE_SIZE" -gt 0 ]; then
    echo "--- 3/5 切成 ${TILE_SIZE}x${TILE_SIZE} 的片 -> $TILED"
    "$PYTHON" scripts/tile_dataset.py \
        --data "$DATASET/data.yaml" --out "$TILED" \
        --size "$TILE_SIZE" --overlap "$TILE_OVERLAP" \
        --min-visibility "$MIN_VISIBILITY" --overwrite
    TRAIN_DATA="$TILED/data.yaml"
    IMGSZ_SET="train.imgsz=$TILE_SIZE"
else
    echo "--- 3/5 跳过切片（TILE_SIZE=0），直接用未切片的数据集"
    TRAIN_DATA="$DATASET/data.yaml"
    IMGSZ_SET=""
fi

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

echo "--- 5/5 切换权重"
# 两个稳定软链接，换模型不用改任何脚本：
#   current.pt —— 最新一次训练的结果，后端与 push_predictions 都读它
#   best.pt    —— 历史上 mAP50 最高的那版，留作回退与对照
ln -sfn "$WORK_DIR/weights/best.pt" "$REPO/runs/detect/current.pt"

map50_of() {   # $1 = 指向某次训练 weights/best.pt 的路径
    local metrics
    metrics="$(dirname "$(dirname "$1")")/metrics.json"
    [ -f "$metrics" ] || return 1
    "$PYTHON" -c "import json,sys; print(json.load(open(sys.argv[1]))[\"map50\"])" "$metrics"
}

NEW_MAP=$(map50_of "$WORK_DIR/weights/best.pt" || echo "")
OLD_BEST=$(readlink -f "$REPO/runs/detect/best.pt" 2>/dev/null || echo "")
OLD_MAP=$([ -n "$OLD_BEST" ] && map50_of "$OLD_BEST" || echo "")

# 数据集每次重训都在变，这个比较只是粗略参考，不是严格的同条件对比
if [ -z "$OLD_MAP" ] || [ -z "$NEW_MAP" ] || \
   awk "BEGIN{exit !($NEW_MAP >= $OLD_MAP)}"; then
    ln -sfn "$WORK_DIR/weights/best.pt" "$REPO/runs/detect/best.pt"
    echo "    best.pt 更新为本次（mAP50 ${NEW_MAP:-未知}，此前 ${OLD_MAP:-无}）"
else
    echo "    best.pt 保持不变（本次 mAP50 $NEW_MAP < 此前 $OLD_MAP）"
fi

echo "--- 归档 -> $MODELS_DIR/film_${SIZE_TAG}_${STAMP}.pt"
ARCHIVE_ARGS=(--weights "$WORK_DIR/weights/best.pt" --out "$MODELS_DIR" --stamp "$STAMP" --overwrite)
[ "$EXPORT_OPENVINO" = "1" ] || ARCHIVE_ARGS+=(--no-openvino)
# 归档失败不该让整轮重训算失败：权重和软链接此时已经就位
"$PYTHON" scripts/export_openvino.py "${ARCHIVE_ARGS[@]}" || echo "    归档失败，跳过（权重仍在 $WORK_DIR/weights/best.pt）"

if [ "$RESTART_BACKEND" = "1" ]; then
    echo "--- 重启 ML backend"
    # 必须显式传 YOLO_WEIGHTS 并清掉继承来的变量：本脚本可能是旧后端的子进程，
    # 它环境里已经 export 了 YOLO_WEIGHTS（旧权重）和可能的 RETRAIN_EXTRA_SET（冒烟参数），
    # 不清掉的话新后端会继续用旧模型、下次训练也会沿用冒烟参数
    # 后端的切片尺寸必须与训练一致，否则输入分布对不上（实测：640 训练的模型
    # 在 160 的片上一个框都出不来）
    env -u RETRAIN_EXTRA_SET YOLO_WEIGHTS="$REPO/runs/detect/current.pt" \
        YOLO_TILE_SIZE="$TILE_SIZE" YOLO_TILE_OVERLAP="$TILE_OVERLAP" \
        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"
echo "归档     -> $MODELS_DIR/film_${SIZE_TAG}_${STAMP}.pt（部署取这份，带戳不会被改写）"
echo "current.pt -> $(readlink "$REPO/runs/detect/current.pt")"
echo "best.pt    -> $(readlink "$REPO/runs/detect/best.pt")"
if [ "$TILE_SIZE" -gt 0 ]; then
    echo "推理请用同样的尺寸：--tile-size $TILE_SIZE --tile-overlap $TILE_OVERLAP"
fi