#!/usr/bin/env bash
# gpu-av1-sample.sh —— 在 GPU 机上给一个源文件跑多档 CRF/preset 的 40 秒样本对比
# 用法： bash gpu-av1-sample.sh <源文件> <输出目录> [线程数]
# 产出： 各档样本 mp4 + 源片段（h264 直拷，可原生播放）+ samples.tsv 汇总 + ssim 日志
set -uo pipefail

SRC="${1:?需要源文件路径}"
OUT="${2:?需要输出目录}"
THREADS="${3:-12}"
SEGS=(900 3000 5200)     # 三个抽样起点（秒）
DUR=40                   # 每段时长（秒）
# 变体：名字|preset|crf|svtav1额外参数
VARIANTS=(
  "p4c19|4|19|enable-tf=0:tune=0"
  "p4c21|4|21|enable-tf=0:tune=0"
  "p4c24|4|24|enable-tf=0:tune=0"
  "p6c21|6|21|enable-tf=0:tune=0"
)

mkdir -p "$OUT"
TSV="$OUT/samples.tsv"
: > "$TSV"
printf 'variant\tpreset\tcrf\tseg\tbytes\tkbps\tssim\n' >> "$TSV"

echo "═══ 样本对比 ═══"
echo "  源：$SRC"
SRC_KBPS=$(ffprobe -v error -select_streams v:0 -show_entries stream=bit_rate -of csv=p=0 "$SRC")
[ -z "$SRC_KBPS" ] || [ "$SRC_KBPS" = "N/A" ] && SRC_KBPS=$(ffprobe -v error -show_entries format=bit_rate -of csv=p=0 "$SRC")
SRC_BYTES_FULL=$(stat -c%s "$SRC")
echo "  源码率约 $((SRC_KBPS/1000)) kbps，整片 $(echo "scale=2;$SRC_BYTES_FULL/1073741824"|bc) GiB"
echo "  抽样：${SEGS[*]} 秒起，各 ${DUR} 秒；线程 $THREADS"
echo

# 源片段（h264 直拷，供对照与网页原生播放）
for s in "${SEGS[@]}"; do
  f="$OUT/src-$s.mp4"
  [ -s "$f" ] && continue
  nice -n 19 ffmpeg -hide_banner -nostdin -loglevel error -y -ss "$s" -t "$DUR" -i "$SRC" \
    -map 0:v:0 -map "0:a:0?" -c:v copy -c:a aac -b:a 128k "$f" 2>/dev/null
  echo "  源片段 $s → $(du -h "$f" | cut -f1)"
done
echo

for v in "${VARIANTS[@]}"; do
  IFS='|' read -r name preset crf sparams <<<"$v"
  for s in "${SEGS[@]}"; do
    out="$OUT/av1-$name-$s.mp4"
    if [ ! -s "$out" ]; then
      t0=$(date +%s)
      nice -n 19 ffmpeg -hide_banner -nostdin -loglevel error -y -ss "$s" -t "$DUR" -i "$SRC" \
        -map 0:v:0 -map "0:a:0?" -c:a aac -b:a 128k \
        -c:v libsvtav1 -preset "$preset" -crf "$crf" \
        -svtav1-params "lp=$THREADS:$sparams" -threads "$THREADS" \
        -pix_fmt yuv420p10le -movflags +faststart "$out" || { echo "  ❌ $name/$s 编码失败"; continue; }
      t1=$(date +%s)
      echo "  编码 $name @ ${s}s → ${t1}-${t0}=$((t1-t0))s，$(du -h "$out" | cut -f1)"
    fi
    # SSIM：同一 -ss 起点解码，帧对齐可靠
    slog="$OUT/ssim-$name-$s.log"
    ssim=$(nice -n 19 ffmpeg -hide_banner -nostdin -loglevel error \
        -ss "$s" -t "$DUR" -i "$SRC" -i "$out" \
        -lavfi "[0:v][1:v]ssim=stats_file=$slog" -f null - 2>&1 | grep -o 'All:[0-9.]*' | head -1 | cut -d: -f2)
    [ -z "$ssim" ] && ssim=$(tail -1 "$slog" 2>/dev/null | grep -o 'All:[0-9.]*' | cut -d: -f2)
    b=$(stat -c%s "$out" 2>/dev/null || echo 0)
    kbps=$(( b * 8 / DUR / 1000 ))
    printf '%s\t%s\t%s\t%s\t%s\t%s\t%s\n' "$name" "$preset" "$crf" "$s" "$b" "$kbps" "${ssim:-?}" >> "$TSV"
  done
done

echo
echo "═══ 汇总（源视频码率 $((SRC_KBPS/1000)) kbps）═══"
column -t -s$'\t' "$TSV"
echo
echo "═══ 各档平均 SSIM 与相对源码率 ═══"
awk -F'\t' -v src="$((SRC_KBPS/1000))" 'NR>1 && $7!="?" {
  n[$1]++; s[$1]+=$7; k[$1]+=$6;
} END {
  for (v in n) printf "  %-8s 平均SSIM %.5f  平均码率 %d kbps（源 %d kbps，%d%%）\n", v, s[v]/n[v], k[v]/n[v], src, k[v]/n[v]*100/src;
}' "$TSV" | sort
echo
echo "✅ 完成，结果在 $OUT"