#!/usr/bin/env bash
set -euo pipefail

# Large CPU expectation benchmarks for two-server runs.
#
# Defaults assume two Intel Xeon Platinum 8558P servers with about 500 GiB RAM
# each.  Override HOSTFILE, PYTHON_BIN, MPIEXEC, or the per-case knobs below as
# needed.

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$ROOT_DIR"

PYTHON_BIN="${PYTHON_BIN:-.venv/bin/python}"
MPIEXEC="${MPIEXEC:-mpiexec}"
HOSTFILE="${HOSTFILE:-hostfile}"

MPS_RANKS="${MPS_RANKS:-8}"
MPS_THREADS="${MPS_THREADS:-12}"
TN_RANKS="${TN_RANKS:-12}"
TN_THREADS="${TN_THREADS:-8}"

export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
export MKL_NUM_THREADS="${MKL_NUM_THREADS:-1}"

run_mpi() {
  local ranks="$1"
  shift
  "$MPIEXEC" -hostfile "$HOSTFILE" -n "$ranks" "$PYTHON_BIN" "$@"
}

run_case() {
  local title="$1"
  shift
  echo
  echo "================================================================================"
  echo "$title"
  echo "================================================================================"
  echo "HOSTFILE=$HOSTFILE PYTHON_BIN=$PYTHON_BIN MPIEXEC=$MPIEXEC"
  echo "OMP_NUM_THREADS=$OMP_NUM_THREADS MKL_NUM_THREADS=$MKL_NUM_THREADS"
  echo "$*"
  "$@"
}

case "${1:-help}" in
  smoke)
    run_case "MPS MPI smoke: n=40 layers=30 bond=2048" \
      run_mpi "$MPS_RANKS" benchmark_cpu_expectation.py \
        --mpi --mps \
        --nqubits "${MPS_SMOKE_NQ:-40}" \
        --nlayers "${MPS_SMOKE_LAYERS:-30}" \
        --bond "${MPS_SMOKE_BOND:-2048}" \
        --torch-threads "$MPS_THREADS" \
        --circuits brickwall_cnot reversed_cnot shifted_cz \
        --observables ring_xz open_zz range2_xx

    run_case "TN MPI smoke: n=32 layers=16 target_slices=12" \
      run_mpi "$TN_RANKS" benchmark_cpu_expectation.py \
        --mpi \
        --nqubits "${TN_SMOKE_NQ:-32}" \
        --nlayers "${TN_SMOKE_LAYERS:-16}" \
        --torch-threads "$TN_THREADS" \
        --circuits brickwall_cnot shifted_cz rxx_rzz \
        --observables ring_xz open_zz range2_xx \
        --tn-target-slices "${TN_SMOKE_SLICES:-12}"
    ;;

  mps-long)
    run_case "MPS MPI long: n=64 layers=48 bond=4096" \
      run_mpi "$MPS_RANKS" benchmark_cpu_expectation.py \
        --mpi --mps \
        --nqubits "${MPS_LONG_NQ:-64}" \
        --nlayers "${MPS_LONG_LAYERS:-48}" \
        --bond "${MPS_LONG_BOND:-4096}" \
        --torch-threads "$MPS_THREADS" \
        --circuits brickwall_cnot reversed_cnot shifted_cz rxx_rzz \
        --observables ring_xz open_zz mixed_local range2_xx
    ;;

  mps-pressure)
    run_case "MPS MPI pressure: n=80 layers=64 bond=4096" \
      run_mpi "$MPS_RANKS" benchmark_cpu_expectation.py \
        --mpi --mps \
        --nqubits "${MPS_PRESSURE_NQ:-80}" \
        --nlayers "${MPS_PRESSURE_LAYERS:-64}" \
        --bond "${MPS_PRESSURE_BOND:-4096}" \
        --torch-threads "$MPS_THREADS" \
        --circuits brickwall_cnot reversed_cnot shifted_cz rxx_rzz swap_scramble \
        --observables ring_xz open_zz mixed_local range2_xx long_z_string
    ;;

  tn-long)
    run_case "TN MPI long: n=36 layers=20 target_slices=24" \
      run_mpi "$TN_RANKS" benchmark_cpu_expectation.py \
        --mpi \
        --nqubits "${TN_LONG_NQ:-36}" \
        --nlayers "${TN_LONG_LAYERS:-20}" \
        --torch-threads "$TN_THREADS" \
        --circuits brickwall_cnot shifted_cz rxx_rzz \
        --observables ring_xz open_zz range2_xx \
        --tn-target-slices "${TN_LONG_SLICES:-24}"
    ;;

  all)
    "$0" smoke
    "$0" mps-long
    "$0" tn-long
    ;;

  help|*)
    cat >&2 <<'EOF'
Usage: tools/run_cpu_large_cases.sh [smoke|mps-long|mps-pressure|tn-long|all]

Common overrides:
  HOSTFILE=hostfile
  PYTHON_BIN=.venv/bin/python
  MPIEXEC=mpiexec
  MPS_RANKS=8 MPS_THREADS=12
  TN_RANKS=12 TN_THREADS=8

Scale overrides:
  MPS_LONG_NQ=64 MPS_LONG_LAYERS=48 MPS_LONG_BOND=4096
  MPS_PRESSURE_NQ=80 MPS_PRESSURE_LAYERS=64 MPS_PRESSURE_BOND=4096
  TN_LONG_NQ=36 TN_LONG_LAYERS=20 TN_LONG_SLICES=24
EOF
    exit 2
    ;;
esac