#!/usr/bin/env bash
# Single-policy trace replay for unified, with tunable overload-factor.
# Used to test direction A: does tightening affinity overflow improve unified?
#
# Usage:
#   OVERLOAD_FACTOR=1.3 bash run_unified_ablation.sh
#   OVERLOAD_FACTOR=1.0 bash run_unified_ablation.sh
#
# Output: $PROJ_DIR/outputs/unified_of${OF}_${DATE}/unified/

set -uo pipefail

PROJ_DIR="${PROJ_DIR:-/home/admin/cpfs/wjh/agentic-kv}"
TRACE="${TRACE:-$PROJ_DIR/traces/w600_r0.0015_st30.jsonl}"
OF="${OVERLOAD_FACTOR:-1.3}"
LMW="${LMETRIC_DECODE_WEIGHT:-0.0}"
TAG_DEFAULT="of${OF/./}"
if [ "$(printf '%s' "$LMW" | grep -v '^0\.\?0*$' || true)" != "" ]; then
    TAG_DEFAULT="${TAG_DEFAULT}_lmw${LMW/./}"
fi
TAG="${TAG:-$TAG_DEFAULT}"
DATE="$(date +%Y%m%d_%H%M)"
OUTROOT="${OUTROOT:-$PROJ_DIR/outputs/unified_${TAG}_${DATE}}"

mkdir -p "$OUTROOT"
echo "=== unified ablation: overload_factor=$OF ==="
echo "Trace : $TRACE"
echo "Out   : $OUTROOT"
echo ""

cleanup() {
    pkill -9 -f cache_aware_proxy 2>/dev/null || true
    pkill -9 -f "vllm serve" 2>/dev/null || true
    pkill -9 -f "EngineCore" 2>/dev/null || true
    sleep 3
}
trap cleanup EXIT
cleanup

cfg_dir="$OUTROOT/unified"
mkdir -p "$cfg_dir"

export EXTRA_PROXY_ARGS="--overload-factor $OF --lmetric-decode-weight $LMW"

echo ""
echo "====== unified ; overload_factor=$OF lmetric_decode_weight=$LMW ======"
bash "$PROJ_DIR/scripts/b3_isolated_policy.sh" "unified" "$TRACE" "$cfg_dir" \
    2>&1 | tee "$cfg_dir/orchestrator.log" | tail -30

pkill -9 -f cache_aware_proxy 2>/dev/null || true
pkill -9 -f "vllm serve" 2>/dev/null || true
pkill -9 -f "EngineCore" 2>/dev/null || true
sleep 5

echo ""
echo "Done. Artifacts: $OUTROOT/unified/metrics.jsonl"