#!/usr/bin/env bash

set -euo pipefail

TP="${1:?usage: run_simulator_group.sh TP MODE}"
MODE="${2:?usage: run_simulator_group.sh TP MODE}"

case "${TP}" in
  1|2|4) ;;
  *) echo "ERROR: TP must be 1, 2, or 4" >&2; exit 2 ;;
esac
case "${MODE}" in
  false|true|both) ;;
  *) echo "ERROR: mode must be false, true, or both" >&2; exit 2 ;;
esac

CHECKOUT="${CHECKOUT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}"
EXPERIMENT_ROOT="${EXPERIMENT_ROOT:-/home/admin/cpfs/wjh/aituner/workload-regime-20260720}"
TRACES_ROOT="${TRACES_ROOT:-${EXPERIMENT_ROOT}/traces}"
FRONTIER_SOURCE="${FRONTIER_SOURCE:-/home/admin/cpfs/wjh/aituner/frontier-t1-dash0-deadc4a}"
REPLAYSERVE_ROOT="${REPLAYSERVE_ROOT:-/home/admin/cpfs/wjh/replayserve}"
PROFILE_ROOT="${PROFILE_ROOT:-/home/admin/cpfs/wjh/aituner/aituner-graph-piecewise-bdc357d/runs/frontier-fidelity-envelope-v1/profiles/profile-v4-trace-final}"
KERNEL_PROFILE_ROOT="${KERNEL_PROFILE_ROOT:-/home/admin/cpfs/wjh/aituner/graph-piecewise-qwen30-20260717/full/frozen-kernel-only}"
ALLREDUCE_CSV="${ALLREDUCE_CSV:-/home/admin/cpfs/wjh/aituner/aituner-graph-piecewise-bdc357d/runs/frontier-fidelity-envelope-v1/profiles/measured-allreduce.csv}"
CUSTOM_PYTHON_DEPS="${CUSTOM_PYTHON_DEPS:-${EXPERIMENT_ROOT}/python-deps-v2}"
BASE_PYTHON_DEPS="${BASE_PYTHON_DEPS:-/home/admin/cpfs/wjh/venvs/qwen36-vllm-0.20.2/lib/python3.12/site-packages}"
SIM_GENERATION="${SIM_GENERATION:-v4}"
PREDICTOR_CACHE_GENERATION="${PREDICTOR_CACHE_GENERATION:-${SIM_GENERATION}}"
RUNNER="${CHECKOUT}/runs/frontier-fidelity-envelope-v1/run_frontier_qwen30_exact_trace_surface.py"

for path in "${TRACES_ROOT}/manifest.json" "${FRONTIER_SOURCE}" \
  "${REPLAYSERVE_ROOT}" "${PROFILE_ROOT}" "${KERNEL_PROFILE_ROOT}" \
  "${ALLREDUCE_CSV}" "${CUSTOM_PYTHON_DEPS}" "${BASE_PYTHON_DEPS}" \
  "${RUNNER}"; do
  [[ -e "${path}" ]] || { echo "ERROR: missing ${path}" >&2; exit 1; }
done

run_mode() {
  local prefix_mode="$1" output_root predictor_cache_root prefix_flag family
  local -a families traces configs command
  output_root="${EXPERIMENT_ROOT}/sim-${SIM_GENERATION}/tp${TP}-prefix-${prefix_mode}"
  predictor_cache_root="${EXPERIMENT_ROOT}/sim-${PREDICTOR_CACHE_GENERATION}/tp${TP}-prefix-${prefix_mode}/cache"
  mkdir -p "${output_root}/provenance"
  [[ -d "${predictor_cache_root}" ]] || {
    echo "ERROR: missing predictor cache ${predictor_cache_root}" >&2
    exit 1
  }

  if [[ "${prefix_mode}" == "true" ]]; then
    prefix_flag="--prefix-caching"
  else
    prefix_flag="--no-prefix-caching"
  fi

  mapfile -t families < <(
    jq -r --argjson prefix "${prefix_mode}" --argjson tp "${TP}" '
      .cases
      | map(select(.prefix_caching == $prefix and ((.tp // $tp) == $tp)))
      | map(.family)
      | unique
      | .[]
    ' "${TRACES_ROOT}/manifest.json"
  )
  [[ "${#families[@]}" -gt 0 ]] || { echo "ERROR: empty family group" >&2; exit 1; }

  for mns in 8 16 32 64; do
    configs+=(--config "tp${TP}_mns${mns}")
  done
  {
    printf 'host=%s\n' "$(hostname)"
    printf 'tp=%s\n' "${TP}"
    printf 'prefix_mode=%s\n' "${prefix_mode}"
    printf 'family_count=%s\n' "${#families[@]}"
    git -C "${CHECKOUT}" rev-parse HEAD
    git -C "${FRONTIER_SOURCE}" rev-parse HEAD
    sha256sum "${TRACES_ROOT}/manifest.json" "${PROFILE_ROOT}/manifest.json" \
      "${KERNEL_PROFILE_ROOT}/manifest.json" "${ALLREDUCE_CSV}"
  } > "${output_root}/provenance/launch.txt"

  for family in "${families[@]}"; do
    traces=()
    command=()
    mapfile -t traces < <(
      jq -r --arg family "${family}" --argjson tp "${TP}" '
        .cases
        | map(select(.family == $family and ((.tp // $tp) == $tp)))
        | sort_by(.global_offered_request_rate)
        | .[]
        | "\(.family)-rho\(.rho | tostring | gsub("\\."; "p"))=\(.public_csv)"
      ' "${TRACES_ROOT}/manifest.json"
    )
    [[ "${#traces[@]}" -eq 5 ]] || {
      echo "ERROR: expected five load anchors for ${family}" >&2
      exit 1
    }
    for trace in "${traces[@]}"; do
      command+=(--trace "${trace}")
    done
    printf 'SIMULATOR_FAMILY_START family=%s tp=%s prefix=%s\n' \
      "${family}" "${TP}" "${prefix_mode}"
    PYTHONPATH="${BASE_PYTHON_DEPS}" /usr/bin/python3 "${RUNNER}" \
      --frontier-source "${FRONTIER_SOURCE}" \
      --replayserve-root "${REPLAYSERVE_ROOT}" \
      --profile-root "${PROFILE_ROOT}" \
      --kernel-profile-root "${KERNEL_PROFILE_ROOT}" \
      --python-deps "${CUSTOM_PYTHON_DEPS}" \
      --output-root "${output_root}" \
      --predictor-cache-root "${predictor_cache_root}" \
      "${command[@]}" "${configs[@]}" \
      --rate-contract trace-window "${prefix_flag}" \
      --cc-backend vidur --allreduce-csv "${ALLREDUCE_CSV}" \
      --timeout-seconds 3600 --predictor-training-job-threads 4 \
      --decode-cuda-graph-mode piecewise --align-real-graph-runtime \
      --store-stage-batch-ledger --resume --continue-on-failure
    cp "${output_root}/frontier_trace_surface_frozen.json" \
      "${output_root}/frontier_trace_surface_${family}.json"
  done
}

if [[ "${MODE}" == "both" ]]; then
  run_mode false
  run_mode true
else
  run_mode "${MODE}"
fi