#!/usr/bin/env bash
# Свип залоченной частоты с замером напряжения и потребления.
# На каждой точке: выставить лок -> дать устояться -> усреднить N замеров.
# По окончании (в том числе по Ctrl-C и при ошибке) исходное состояние восстанавливается.
#
#   sudo ./sweep.sh -i 0 --from 900 --to 1500 --step 50
#   sudo ./sweep.sh -i 0 --mem --from 5000 --to 9500 --step 250 --dwell 20
#   sudo ./sweep.sh -i 0 --from 900 --to 1400 --step 50 --hashrate-cmd 'my-miner-api-hashrate 0'
#
# ВНИМАНИЕ: скрипт МЕНЯЕТ состояние карты (лочит частоты). Запускать под нагрузкой --
# в простое цифры бессмысленны, карта сидит в idle-pstate.

set -uo pipefail

NVTOOL="${NVTOOL:-/hive/sbin/nvtool}"
IDX=0; FROM=""; TO=""; STEP=50; DWELL=15; SAMPLES=5; DOMAIN=core
HASH_CMD=""; OUT=""; ASSUME_YES=0

die() { echo "ОШИБКА: $*" >&2; exit 1; }

usage() {
	sed -n '2,14p' "$0" | sed 's/^# \?//'
	cat <<-EOF

	  -i N              индекс GPU (по умолчанию 0)
	  --from MHz        начало диапазона (обязательно)
	  --to MHz          конец диапазона (обязательно)
	  --step MHz        шаг (по умолчанию 50)
	  --mem             свипать память вместо ядра
	  --dwell SEC       пауза на устаканивание после смены лока (по умолчанию 15)
	  --samples N       замеров на точку (по умолчанию 5, раз в секунду)
	  --hashrate-cmd C  команда, печатающая хешрейт числом -- попадёт в таблицу
	  --csv FILE        куда сохранить таблицу (по умолчанию sweep-<домен>-<дата>.csv)
	  -y                не спрашивать подтверждения
	EOF
	exit 0
}

while [ $# -gt 0 ]; do
	case "$1" in
		-i) IDX="$2"; shift 2 ;;
		--from) FROM="$2"; shift 2 ;;
		--to) TO="$2"; shift 2 ;;
		--step) STEP="$2"; shift 2 ;;
		--mem) DOMAIN=mem; shift ;;
		--dwell) DWELL="$2"; shift 2 ;;
		--samples) SAMPLES="$2"; shift 2 ;;
		--hashrate-cmd) HASH_CMD="$2"; shift 2 ;;
		--csv) OUT="$2"; shift 2 ;;
		-y|--yes) ASSUME_YES=1; shift ;;
		-h|--help) usage ;;
		*) die "неизвестный аргумент: $1 (--help)" ;;
	esac
done

[ -n "$FROM" ] && [ -n "$TO" ] || die "нужны --from и --to"
[ "$STEP" -gt 0 ] 2>/dev/null || die "--step должен быть больше нуля"
[ -x "$NVTOOL" ] || die "нет $NVTOOL (задайте NVTOOL=/путь/к/nvtool)"
[ "$(id -u)" = 0 ] || die "нужен root"
OUT="${OUT:-sweep-$DOMAIN-$(date +%Y%m%d-%H%M%S).csv}"

if [ "$DOMAIN" = core ]; then SET_KEY=--setcore; GET_LABEL="GPU CLOCKS CURRENT"
else                          SET_KEY=--setmem;  GET_LABEL="MEM CLOCKS CURRENT"; fi

nv() { "$NVTOOL" --index "$IDX" "$@" 2>&1; }

# "  GPU CLOCKS CURRENT: 1234 MHz" -> 1234
field() { sed -n "s/.*$1:[[:space:]]*\([0-9.]\+\).*/\1/p" <<<"$2" | head -1; }

# 71599732246 -> "71.60 G"; в CSV при этом уходит исходное число
human() { awk -v v="$1" 'BEGIN{
	if (v >= 1e9) printf "%.2f G", v/1e9;
	else if (v >= 1e6) printf "%.2f M", v/1e6;
	else if (v >= 1e3) printf "%.2f K", v/1e3;
	else printf "%.2f", v }'; }

# --- исходное состояние, чтобы вернуть его как было ---
ORIG=$(nv --clocksinfo)
ORIG_CORE=$(field "GPU CLOCKS LOCKED" "$ORIG"); ORIG_MEM=$(field "MEM CLOCKS LOCKED" "$ORIG")
ORIG_CORE=${ORIG_CORE:-0}; ORIG_MEM=${ORIG_MEM:-0}

restore() {
	echo
	echo "восстанавливаю: lock core=$ORIG_CORE mem=$ORIG_MEM"
	nv --setcore "$ORIG_CORE" >/dev/null
	nv --setmem  "$ORIG_MEM"  >/dev/null
	echo "таблица: $OUT"
}
# restore висит на EXIT и потому срабатывает ровно один раз; обработчик сигнала
# только выходит -- иначе после Ctrl-C свип продолжил бы крутить частоты дальше
on_signal() { echo; echo "прервано пользователем"; exit 130; }
trap restore EXIT
trap on_signal INT TERM

USAGE_LINE=$(nv --usage)
GPU_UTIL=$(field "GPU USAGE" "$USAGE_LINE"); GPU_UTIL=${GPU_UTIL:-0}
echo "GPU $IDX, домен $DOMAIN, диапазон $FROM..$TO шаг $STEP, загрузка сейчас ${GPU_UTIL}%"
[ "${GPU_UTIL%%.*}" -lt 50 ] 2>/dev/null && \
	echo "ВНИМАНИЕ: загрузка ниже 50% -- в простое замеры бессмысленны, запускайте под нагрузкой"
echo "текущий лок: core=$ORIG_CORE mem=$ORIG_MEM (он же будет восстановлен)"
PL_LINE=$(nv --curpl --defpl)
echo "лимит мощности: $(sed -n 's/.*POWER LIMIT CURRENT:[[:space:]]*//p' <<<"$PL_LINE" | head -1)" \
     "(по умолчанию $(sed -n 's/.*POWER LIMIT DEFAULT:[[:space:]]*//p' <<<"$PL_LINE" | head -1))"

if [ "$ASSUME_YES" != 1 ]; then
	printf "менять частоты на этой карте? [y/N] "
	read -r a; [ "$a" = y ] || [ "$a" = Y ] || { trap - EXIT; echo "отменено"; exit 0; }
fi

HDR="set_mhz,actual_mhz,mv,watt,mhz_per_watt"
[ -n "$HASH_CMD" ] && HDR="$HDR,hashrate,hash_per_watt"
HDR="$HDR,throttle"
echo "$HDR" > "$OUT"

echo
if [ -n "$HASH_CMD" ]; then
	echo "     set     actual    mV        W        MHz/W    hash        hash/W      throttle"
	echo "  ----------------------------------------------------------------------------------------"
else
	echo "     set     actual    mV        W        MHz/W    throttle"
	echo "  ------------------------------------------------------------------"
fi
CAPPED=0

for ((f = FROM; f <= TO; f += STEP)); do
	nv "$SET_KEY" "$f" >/dev/null
	sleep "$DWELL"

	sum_c=0; sum_v=0; sum_w=0; n=0
	for ((s = 0; s < SAMPLES; s++)); do
		out=$(nv --clocks --vcore --power)
		c=$(field "$GET_LABEL" "$out"); v=$(field "CORE VOLTAGE" "$out"); w=$(field "POWER USAGE" "$out")
		[ -n "$c" ] && [ -n "$w" ] || continue
		sum_c=$(awk -v a="$sum_c" -v b="$c" 'BEGIN{print a+b}')
		sum_v=$(awk -v a="$sum_v" -v b="${v:-0}" 'BEGIN{print a+b}')
		sum_w=$(awk -v a="$sum_w" -v b="$w" 'BEGIN{print a+b}')
		n=$((n + 1))
		sleep 1
	done
	[ "$n" -gt 0 ] || { echo "точка $f: замеры не сняты, пропуск"; continue; }

	read -r ac av aw eff <<<"$(awk -v c="$sum_c" -v v="$sum_v" -v w="$sum_w" -v n="$n" \
		'BEGIN{c/=n; v/=n; w/=n; printf "%.0f %.1f %.1f %.2f", c, v, w, (w>0 ? c/w : 0)}')"

	thr=$(nv --throttle | sed -n 's/.*THROTTLE:[[:space:]]*//p' | head -1)
	thr=${thr:-none}

	# если фактическая частота ниже заданной -- точка упёрлась в другой кап,
	# и это уже не свип частоты, а замер того же режима под разными метками
	mark=" "
	if awk -v a="$ac" -v f="$f" 'BEGIN{exit !(a < f*0.98)}'; then mark="!"; CAPPED=$((CAPPED + 1)); fi

	if [ -n "$HASH_CMD" ]; then
		h=$(eval "$HASH_CMD" 2>/dev/null | tr -dc '0-9.' ); h=${h:-0}
		hpw=$(awk -v h="$h" -v w="$aw" 'BEGIN{printf "%.4f", (w>0 ? h/w : 0)}')
		printf "  %s %-7s %-9s %-9s %-8s %-8s %-11s %-11s %s\n" \
			"$mark" "$f" "$ac" "$av" "$aw" "$eff" "$(human "$h")" "$(human "$hpw")" "$thr"
		echo "$f,$ac,$av,$aw,$eff,$h,$hpw,\"$thr\"" >> "$OUT"
	else
		printf "  %s %-7s %-9s %-9s %-8s %-8s %s\n" "$mark" "$f" "$ac" "$av" "$aw" "$eff" "$thr"
		echo "$f,$ac,$av,$aw,$eff,\"$thr\"" >> "$OUT"
	fi
done

if [ "${CAPPED:-0}" -gt 0 ]; then
	echo
	echo "ВНИМАНИЕ: точек, помеченных '!': $CAPPED -- карта не добрала заданную частоту."
	echo "Смотрите колонку throttle: PowerLimit -- упёрлись в лимит мощности (--setpl),"
	echo "HWThermal/SWThermal -- в температуру. Выше этой точки свип частоты бессмыслен:"
	echo "меняется метка, а режим один и тот же. Поднимите лимит или свипайте сам лимит."
fi