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

NQ="${NQ:-34}"
LAYERS="${LAYERS:-20}"
BOND="${BOND:-512}"
SEED="${SEED:-42}"
RANKS="${RANKS:-1 2 4}"
THREADS="${THREADS:-32 32 16}"
PYTHON_BIN="${PYTHON_BIN:-.venv/bin/python}"
MPIEXEC="${MPIEXEC:-mpiexec}"
CIRCUIT="${CIRCUIT:-brickwall_cnot}"
OBSERVABLE="${OBSERVABLE:-ring_xz}"
EXACT="${EXACT:-0}"

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

if [[ "${1:-help}" != "run" ]]; then
  cat >&2 <<'EOF'
Usage: tools/run_vidal_segment_mpi_scan.sh run

Overrides:
  NQ=34 LAYERS=20 BOND=512 SEED=42
  RANKS="1 2 4" THREADS="32 32 16"
  CIRCUIT=brickwall_cnot OBSERVABLE=ring_xz
  EXACT=1
  PYTHON_BIN=.venv/bin/python MPIEXEC=mpiexec
EOF
  if [[ "${1:-help}" == "help" ]]; then
    exit 0
  fi
  exit 2
fi

read -r -a ranks <<< "$RANKS"
read -r -a threads <<< "$THREADS"

if [[ "${#ranks[@]}" != "${#threads[@]}" ]]; then
  echo "RANKS and THREADS must have the same number of entries." >&2
  exit 2
fi

common=(
  --nqubits "$NQ"
  --nlayers "$LAYERS"
  --bond "$BOND"
  --seed "$SEED"
  --mps
  --circuits "$CIRCUIT"
  --observables "$OBSERVABLE"
)

if [[ "$EXACT" == "1" ]]; then
  common+=(--exact)
fi

for idx in "${!ranks[@]}"; do
  nrank="${ranks[$idx]}"
  nthr="${threads[$idx]}"
  if [[ "$nrank" == "1" ]]; then
    echo "== Vidal serial ranks=1 torch_threads=$nthr =="
    "$PYTHON_BIN" -u benchmark_cpu_expectation.py \
      "${common[@]}" --torch-threads "$nthr"
  else
    echo "== Vidal segmented MPI ranks=$nrank torch_threads=$nthr =="
    "$MPIEXEC" -n "$nrank" "$PYTHON_BIN" -u benchmark_cpu_expectation.py \
      "${common[@]}" --torch-threads "$nthr" --mpi
  fi
done