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

# Focused Vidal/MPS expectation test cases for 1D chain circuits.
#
# These cases intentionally avoid qmatchatea and generic TN paths.  They target
# the current supported scope: one-qubit gates, adjacent two-qubit gates, and
# Pauli-sum expectation values on a 1D chain.

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

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

THREADS="${THREADS:-32}"
MPI_RANKS="${MPI_RANKS:-16}"
MPI_THREADS="${MPI_THREADS:-12}"

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

run() {
  echo
  echo "--------------------------------------------------------------------------------"
  echo "$*"
  echo "--------------------------------------------------------------------------------"
  "$@"
}

case "${1:-help}" in
  smoke)
    # Short correctness-oriented run.  Useful before starting long jobs.
    run "$PYTHON_BIN" -u benchmark_cpu_expectation.py \
      --mps \
      --nqubits 40 \
      --nlayers 10 \
      --bond 2048 \
      --torch-threads "$THREADS" \
      --circuits brickwall_cnot reversed_cnot shifted_cz rxx_rzz \
      --observables ring_xz open_zz range2_xx long_z_string
    ;;

  convergence)
    # Same circuit/observable, increasing bond.  Check value convergence.
    for bond in ${BONDS:-4096 16384 65536}; do
      run "$PYTHON_BIN" -u benchmark_cpu_expectation.py \
        --mps \
        --nqubits "${NQ:-80}" \
        --nlayers "${LAYERS:-16}" \
        --bond "$bond" \
        --torch-threads "$THREADS" \
        --circuits "${CIRCUIT:-brickwall_cnot}" \
        --observables "${OBSERVABLE:-ring_xz}"
    done
    ;;

  single-long)
    # Single long Vidal run.  On node-3, a similar n=40,l=30,bond=2048 case
    # took about 9 minutes for one expectation.  This one is meant to be longer.
    run "$PYTHON_BIN" -u benchmark_cpu_expectation.py \
      --mps \
      --nqubits "${NQ:-80}" \
      --nlayers "${LAYERS:-16}" \
      --bond "${BOND:-65536}" \
      --torch-threads "$THREADS" \
      --circuits "${CIRCUIT:-brickwall_cnot}" \
      --observables "${OBSERVABLE:-ring_xz}"
    ;;

  suite-long)
    # Application-style multi-circuit, multi-observable MPS run.
    # This is intentionally multi-term and should run much longer than single-long.
    run "$PYTHON_BIN" -u benchmark_cpu_expectation.py \
      --mps \
      --nqubits "${NQ:-80}" \
      --nlayers "${LAYERS:-16}" \
      --bond "${BOND:-65536}" \
      --torch-threads "$THREADS" \
      --circuits brickwall_cnot reversed_cnot shifted_cz rxx_rzz \
      --observables ring_xz open_zz mixed_local range2_xx long_z_string
    ;;

  mpi-long)
    # Multi-node Vidal segmented MPS run.  Uses HOSTFILE.
    run "$MPIEXEC" -hostfile "$HOSTFILE" -n "$MPI_RANKS" "$PYTHON_BIN" -u benchmark_cpu_expectation.py \
      --mpi --mps \
      --nqubits "${NQ:-80}" \
      --nlayers "${LAYERS:-16}" \
      --bond "${BOND:-65536}" \
      --torch-threads "$MPI_THREADS" \
      --circuits brickwall_cnot reversed_cnot shifted_cz rxx_rzz \
      --observables ring_xz open_zz mixed_local range2_xx long_z_string
    ;;

  stress)
    # Heavier entanglement.  Start only after single-long is stable.
    run "$PYTHON_BIN" -u benchmark_cpu_expectation.py \
      --mps \
      --nqubits "${NQ:-80}" \
      --nlayers "${LAYERS:-18}" \
      --bond "${BOND:-262144}" \
      --torch-threads "${THREADS:-48}" \
      --circuits "${CIRCUIT:-rxx_rzz}" \
      --observables ring_xz open_zz range2_xx
    ;;

  help|*)
    cat <<'EOF'
Usage: ./run_vidal_mps_cases.sh [smoke|convergence|single-long|suite-long|mpi-long|stress]

Common overrides:
  PYTHON_BIN=.venv/bin/python
  THREADS=32
  OMP_NUM_THREADS=1 MKL_NUM_THREADS=1

Single-node scale overrides:
  NQ=80 LAYERS=16 BOND=65536
  CIRCUIT=brickwall_cnot
  OBSERVABLE=ring_xz
  BONDS="4096 16384 65536"   # for convergence mode

Multi-node overrides:
  HOSTFILE=hostfile
  MPI_RANKS=16 MPI_THREADS=12

Recommended first runs:
  ./run_vidal_mps_cases.sh smoke
  ./run_vidal_mps_cases.sh convergence
  ./run_vidal_mps_cases.sh single-long
EOF
    ;;
esac