#!/usr/bin/env bash

set -euo pipefail

BASE_CAMPAIGN_ROOT="${BASE_CAMPAIGN_ROOT:?BASE_CAMPAIGN_ROOT is required}"
OUTPUT_ROOT="${OUTPUT_ROOT:?OUTPUT_ROOT is required}"
ALLREDUCE_CSV="${ALLREDUCE_CSV:?ALLREDUCE_CSV is required}"
RUNNER_DIR="${RUNNER_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)}"
PROFILE_ROOT="${PROFILE_ROOT:-/home/admin/cpfs/wjh/aituner/qwen235-v020-profiles-20260719-r1/frozen}"
FRONTIER_SOURCE="${FRONTIER_SOURCE:-/home/admin/cpfs/wjh/aituner/frontier-q235-v020-5b953f5}"
REPLAYSERVE_ROOT="${REPLAYSERVE_ROOT:-/home/admin/cpfs/wjh/replayserve}"
VENV_ROOT="${VENV_ROOT:-/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1}"
PYTHON_DEPS="${PYTHON_DEPS:-${VENV_ROOT}/lib/python3.12/site-packages}"

mkdir -p "${OUTPUT_ROOT}"/{provenance,sim,analysis}
exec > >(tee -a "${OUTPUT_ROOT}/controller.log") 2>&1
echo "Q235_COLLECTIVE_ABLATION_LAUNCH_ECHO host=$(hostname) base=${BASE_CAMPAIGN_ROOT} simulator=Frontier-6e8e0d8 configs={TP4/EP1,TP8/EP8}xMNS{64,128} cases={fixed-pd,fixed-po,trace-pd,trace-po} real_rerun=false measured_collective=${ALLREDUCE_CSV} analytical_fallback=forbidden expected_wall=1-6h output=${OUTPUT_ROOT}"
date -u +START_UTC=%Y-%m-%dT%H:%M:%SZ
test "$(git -C "${FRONTIER_SOURCE}" rev-parse HEAD)" = "6e8e0d845bceff11b0b62cb29df3a1a93411fdd4"
sha256sum "${BASH_SOURCE[0]}" "${ALLREDUCE_CSV}" \
  "${PROFILE_ROOT}/manifest.json" \
  "${BASE_CAMPAIGN_ROOT}/provenance/runtime-contract.json" \
  > "${OUTPUT_ROOT}/provenance/input.sha256"
git -C "${RUNNER_DIR}" rev-parse HEAD > "${OUTPUT_ROOT}/provenance/aituner.commit"
git -C "${FRONTIER_SOURCE}" rev-parse HEAD > "${OUTPUT_ROOT}/provenance/frontier.commit"

run_sim() {
  local case_name="$1" prefix_flag="$2"
  "${VENV_ROOT}/bin/python" "${RUNNER_DIR}/run_frontier_qwen235_v020_surface.py" \
    --frontier-source "${FRONTIER_SOURCE}" --replayserve-root "${REPLAYSERVE_ROOT}" \
    --profile-root "${PROFILE_ROOT}" --python-deps "${PYTHON_DEPS}" \
    --output-root "${OUTPUT_ROOT}/sim/${case_name}" \
    --runtime-contract "${BASE_CAMPAIGN_ROOT}/provenance/runtime-contract.json" \
    --trace-tp "4=${BASE_CAMPAIGN_ROOT}/traces/${case_name}/tp4/public/frontier.csv" \
    --trace-tp "8=${BASE_CAMPAIGN_ROOT}/traces/${case_name}/tp8/public/frontier.csv" \
    "${prefix_flag}" --allreduce-csv "${ALLREDUCE_CSV}" --resume
}

run_sim fixed-pd --no-prefix-caching
run_sim fixed-po --no-prefix-caching
run_sim trace-pd --prefix-caching
run_sim trace-po --prefix-caching

"${VENV_ROOT}/bin/python" "${RUNNER_DIR}/analyze_qwen235_v020_campaign.py" \
  --campaign-root "${BASE_CAMPAIGN_ROOT}" --sim-root "${OUTPUT_ROOT}" \
  --json-output "${OUTPUT_ROOT}/analysis/comparison.json" \
  --markdown-output "${OUTPUT_ROOT}/analysis/comparison.md"
find "${OUTPUT_ROOT}" -type f ! -path '*/provenance/artifacts.sha256' -print0 \
  | sort -z | xargs -0 sha256sum > "${OUTPUT_ROOT}/provenance/artifacts.sha256"
date -u +END_UTC=%Y-%m-%dT%H:%M:%SZ
echo Q235_COLLECTIVE_ABLATION_COMPLETE