#!/bin/bash
# 多轮 A/B：同一对比在 N 轮独立进程里各跑一次，逐 shape 取轮次中位。
#
# 为什么需要它：单对观测无法说明自己稳不稳。S5000 实测 topk 消除了 32 个
# convert_layout，42 个 shape 里却有 17 个在两次同等有效的运行之间改变了 delta
# 的符号——那不是计时口径的问题（图捕获与 direct 的延迟比是 1.00x），是收益本身
# 落在零附近、单次测量给不出方向。多轮之后才谈得上"这个判定成立"。
#
# 口径：第 1 轮 A 侧 record autotune config，此后每一次运行（含 A 侧自己的后续
# 轮次）都 REPLAY_FROM 那一份。config 因此跨轮跨侧全部锁死，剩下的变量只有被测
# 开关与测量噪声。逐轮交替先后顺序，抵消单向漂移。
#
# 用法：
#   SWITCH=FLAGTREE_MUSA_RLC_ENHANCE OPS="topk masked_select" ROUNDS=6 \
#     ZL_BACKEND=musa FLAGGEMS_DIR=/path/to/FlagGems bash ab_repeat.sh
#
#   OPS 元素为 OP 或 OP:OP_FILE（省略时 OP_FILE=OP）。
#   SWITCH 是被测环境变量，A 侧设 0、B 侧设 1。
#
# 对比轴有效性由脚本自己查：第 1 轮 A/B 跑完后比两侧的 ttgir 指纹，完全一致说明
# 这个开关在这个算子上没改变任何生成代码——继续跑满 N 轮只会得到一个"无差异"的假
# 结论（同 ab_fold_test.sh 的告诫）。这类算子默认跳过；开关确实只影响运行时而不影响
# IR 时，用 SWITCH_ALLOW_NO_IR_DIFF=1 放行。
#
# 产物：runs/ab_repeat_<开关>_<时间戳>/<op>/r<k>_<A|B>/，跑完对每个算子调用
# batch_summary.py --ab-rounds 出报告。
set -uo pipefail

cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

SWITCH_ALLOW_NO_IR_DIFF="${SWITCH_ALLOW_NO_IR_DIFF:-0}"

# 一个 run 生成代码的指纹：ttgir 相对路径 + 内容。路径进指纹是因为"变体集合变了"
# 与"某个变体内容变了"同样算改变，而只 cat 内容会把前者洗掉。
ir_fingerprint() {
  local dir="$1/ttgir"
  [[ -d "$dir" ]] || return 1
  local out
  out=$( { find "$dir" -type f | LC_ALL=C sort | while IFS= read -r f; do
             printf '%s
' "${f#"$dir"/}"; cat "$f"
           done; } | sha256sum | cut -d' ' -f1 )
  [[ -n "$out" ]] && printf '%s' "$out"
}

: "${SWITCH:?SWITCH is required, e.g. FLAGTREE_MUSA_RLC_ENHANCE}"
OPS="${OPS:-}"
: "${OPS:?OPS is required, e.g. \"topk masked_select\"}"
ROUNDS="${ROUNDS:-6}"

if [[ -t 1 && -z "${NO_COLOR:-}" ]]; then
  export FLAGGEMS_PERF_COLOR=always
fi

ROOT="runs/ab_repeat_${SWITCH}_$(date +%Y%m%d-%H%M%S)"
mkdir -p "$ROOT"
LOG="$ROOT/summary.log"

{
  echo ">>> switch=$SWITCH  rounds=$ROUNDS  ops=$OPS"
  for entry in $OPS; do
    op="${entry%%:*}"
    op_file="${entry#*:}"
    [[ "$op_file" == "$entry" ]] && op_file="$op"
    out="$ROOT/$op"
    mkdir -p "$out"

    # 第 1 轮 A 侧 record，其余全部回放它，config 跨轮跨侧一致
    echo "########## $op | round 1 | A ($SWITCH=0, record) ##########"
    if ! OP="$op" OP_FILE="$op_file" OUT_DIR="$out/r1_A" \
         env "$SWITCH=0" bash run_pytest.sh; then
      echo "!!! FAILED $op round 1 A; skipping this operator"
      continue
    fi
    REC="$out/r1_A"

    ok=1; reason="有失败轮次"
    for ((k = 1; k <= ROUNDS; k++)); do
      # 奇数轮先 A 后 B，偶数轮反过来
      if (( k % 2 == 1 )); then order="A B"; else order="B A"; fi
      for side in $order; do
        [[ $k -eq 1 && $side == A ]] && continue   # 已经跑过
        val=$([[ $side == A ]] && echo 0 || echo 1)
        echo "########## $op | round $k | $side ($SWITCH=$val, replay) ##########"
        if ! OP="$op" OP_FILE="$op_file" OUT_DIR="$out/r${k}_${side}" \
             REPLAY_FROM="$REC" env "$SWITCH=$val" bash run_pytest.sh; then
          echo "!!! FAILED $op round $k $side"
          ok=0
        fi
      done

      if (( k == 1 && ok )); then
        fa=$(ir_fingerprint "$out/r1_A") || fa=""
        fb=$(ir_fingerprint "$out/r1_B") || fb=""
        if [[ -z "$fa" || -z "$fb" ]]; then
          echo "??? $op: 没有 ttgir 产物，跳过对比轴有效性检查"
        elif [[ "$fa" == "$fb" ]]; then
          echo "!!! SWITCH_NO_IR_DIFF op=$op switch=$SWITCH"
          echo "!!! 两侧 ttgir 逐字节相同：这个开关没有改变 $op 的生成代码。"
          echo "!!! 继续跑 $ROUNDS 轮只会测出噪声，并报成一个'无差异'结论。"
          echo "!!! 先确认开关名拼写与它在被测源码里的读取点；确属只影响运行时的"
          echo "!!! 开关，用 SWITCH_ALLOW_NO_IR_DIFF=1 重跑。"
          if [[ "$SWITCH_ALLOW_NO_IR_DIFF" != "1" ]]; then
            ok=0; reason="开关未改变该算子的生成代码"
            break
          fi
        else
          echo ">>> 对比轴有效：$op 两侧 ttgir 不同（A=${fa:0:12} B=${fb:0:12}）"
        fi
      fi
    done

    echo "########## $op | report ##########"
    if (( ok )); then
      python batch_summary.py --ab-rounds "$out" --label "$op" \
        || echo "!!! report failed: $op"
    else
      # 缺轮的算子照样出报告会让人误以为是完整的 N 轮
      echo "!!! $op $reason，不出报告；产物在 $out"
    fi
  done
  echo "########## ALL DONE ##########"
} 2>&1 | tee "$LOG"

sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG"

echo ">>> 总 log: $(pwd)/$LOG"
echo ">>> 逐算子重出报告: python batch_summary.py --ab-rounds $ROOT/<op> --label <op>"
echo ">>> 解读要点: 逐轮区间宽说明判定不稳；单侧噪声大说明机器不稳；两者含义不同"