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

# ------------------------------
# 资源限制配置（cgroup v2 + systemd-run）
# ------------------------------
# 说明：
# 1) 这个脚本会把训练进程放进一个临时的 systemd scope 中，并施加 CPU/内存上限。
# 2) 仅限制“本次训练进程”，不会永久改系统配置。
# 3) 下面变量都支持“环境变量覆盖”，即你可以在命令前临时指定。
#
# CPU 核数基准（默认 20）：
# 例如你的机器按 20 核预算来算，可保持默认。
CORES="${CORES:-20}"
# CPU 占用百分比（默认 100）：
# 最终会与 CORES 相乘得到 CPUQuota。
# 例：CORES=20, UTIL_PERCENT=100 -> CPUQuota=2000%（约 20 核等效）
UTIL_PERCENT="${UTIL_PERCENT:-100}"
CPU_QUOTA="$((CORES * UTIL_PERCENT))%"

# 内存上限（默认关闭）：
# 设为具体值（如 24G/28G）可限制训练最多占用内存；
# 设为 0/none/off/infinity 表示不设置 cgroup 内存上限。
MEMORY_MAX="${MEMORY_MAX:-0}"

# ------------------------------
# 精简运行参数（由 scripts/gyms/Walk.py 读取）
# ------------------------------
# 仅保留最常用开关，避免超长环境变量命令。
GYM_CPU_MODE="${GYM_CPU_MODE:-train}"
GYM_CPU_TRAIN_STAGE="${GYM_CPU_TRAIN_STAGE:-walk}"
GYM_CPU_TRAIN_MODEL="${GYM_CPU_TRAIN_MODEL:-}"
GYM_CPU_TEST_MODEL="${GYM_CPU_TEST_MODEL:-scripts/gyms/logs/Walk_R0_004/best_model.zip}"
GYM_CPU_TEST_FOLDER="${GYM_CPU_TEST_FOLDER:-scripts/gyms/logs/Walk_R0_004/}"

# Python 解释器选择策略：
# 1) 优先使用你手动传入的 PYTHON_BIN
# 2) 其次用当前激活 conda 环境（CONDA_PREFIX/bin/python）
# 3) 再回退到默认 mujoco 环境路径
# 4) 最后尝试系统 python / python3
DEFAULT_PYTHON="/home/solren/Downloads/Anaconda/envs/mujoco/bin/python"
CONDA_PYTHON="${CONDA_PREFIX:-}/bin/python"

# 安全保护：不要用 sudo 运行。
# 原因：sudo 可能导致 conda 环境与用户会话环境不一致，
# 会引发 python 路径丢失、systemd --user 会话不可见等问题。
if [[ "${EUID}" -eq 0 ]]; then
	echo "Do not run this script with sudo; run as your normal user in conda env 'mujoco'."
	exit 1
fi

# 解析最终使用的 Python 可执行文件。
if [[ -n "${PYTHON_BIN:-}" ]]; then
	PYTHON_EXEC="${PYTHON_BIN}"
elif [[ -n "${CONDA_PREFIX:-}" && -x "${CONDA_PYTHON}" ]]; then
	PYTHON_EXEC="${CONDA_PYTHON}"
elif [[ -x "${DEFAULT_PYTHON}" ]]; then
	PYTHON_EXEC="${DEFAULT_PYTHON}"
elif command -v python >/dev/null 2>&1; then
	PYTHON_EXEC="$(command -v python)"
elif command -v python3 >/dev/null 2>&1; then
	PYTHON_EXEC="$(command -v python3)"
else
	echo "No Python executable found. Set PYTHON_BIN=/abs/path/to/python and retry."
	exit 1
fi

# 脚本所在目录（绝对路径），便于后续定位模块/相对路径。
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"

# 打印当前生效配置，方便排障和复现实验。
echo "Starting training with limits: CPU=${CPU_QUOTA}, Memory=${MEMORY_MAX}"
echo "Mode: ${GYM_CPU_MODE}"
echo "Run knobs: GYM_CPU_MODE=${GYM_CPU_MODE}, GYM_CPU_TRAIN_STAGE=${GYM_CPU_TRAIN_STAGE}"
echo "Using Python: ${PYTHON_EXEC}"
if [[ -n "${CONDA_DEFAULT_ENV:-}" ]]; then
	echo "Detected conda env: ${CONDA_DEFAULT_ENV}"
fi

SYSTEMD_PROPS=("-p" "CPUQuota=${CPU_QUOTA}")
case "${MEMORY_MAX,,}" in
	0|none|off|infinity)
		echo "MemoryMax is disabled for this run (no cgroup memory cap)."
		;;
	*)
		SYSTEMD_PROPS+=("-p" "MemoryMax=${MEMORY_MAX}")
		;;
esac

# 使用 systemd-run --user --scope 启动“受限资源”的训练进程：
# - CPUQuota: 总 CPU 配额
# - MemoryMax: 最大内存
# - env ...  : 显式传递训练参数到 Python 进程
# - python -m scripts.gyms.Walk: 以模块方式启动训练入口
systemd-run --user --scope \
	"${SYSTEMD_PROPS[@]}" \
	env \
	GYM_CPU_MODE="${GYM_CPU_MODE}" \
	GYM_CPU_TRAIN_STAGE="${GYM_CPU_TRAIN_STAGE}" \
	GYM_CPU_TRAIN_MODEL="${GYM_CPU_TRAIN_MODEL}" \
	GYM_CPU_TEST_MODEL="${GYM_CPU_TEST_MODEL}" \
	GYM_CPU_TEST_FOLDER="${GYM_CPU_TEST_FOLDER}" \
	"${PYTHON_EXEC}" "-m" "scripts.gyms.Walk"