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

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

CASE="${CASE:-main1}"
OBSERVABLES="${OBSERVABLES:-long_z_string}"
NQUBITS="${NQUBITS:-34}"
NLAYERS="${NLAYERS:-20}"
TORCH_THREADS="${TORCH_THREADS:-48}"
SEARCH_REPEATS="${SEARCH_REPEATS:-2048}"
SEARCH_TIME="${SEARCH_TIME:-300}"
TN_TARGET_SIZE="${TN_TARGET_SIZE:-8589934592}"
TN_TARGET_SLICES="${TN_TARGET_SLICES:-}"

PYTHON_BIN="${PYTHON_BIN:-.venv/bin/python}"
DTYPE="${DTYPE:-complex64}"
TREE_DIR="${TREE_DIR:-trees/contest_tn}"
DASK_ADDRESS="${DASK_ADDRESS:-tcp://10.20.1.103:8786}"
MPIEXEC_FULL="${MPIEXEC_FULL:-mpirun -np 4 -hostfile /home/yx/qibotn/hostfile -perhost 2}"
SYNC_TREES="${SYNC_TREES:-1}"
SYNC_HOSTS="${SYNC_HOSTS:-${WORKER_HOSTS:-}}"
SSH_BIN="${SSH_BIN:-ssh}"

export TCM_ENABLE="${TCM_ENABLE:-1}"

tn_slice_args=(--tn-target-size "$TN_TARGET_SIZE")
if [[ -n "$TN_TARGET_SLICES" ]]; then
  tn_slice_args+=(--tn-target-slices "$TN_TARGET_SLICES")
fi

is_local_host() {
  local host="$1"
  [[ "$host" == "localhost" || "$host" == "127.0.0.1" ]] && return 0
  [[ "$host" == "$(hostname)" ]] && return 0
  [[ "$host" == "$(hostname -f 2>/dev/null || true)" ]] && return 0
  hostname -I 2>/dev/null | tr ' ' '\n' | grep -qx "$host"
}

sync_trees_to_hosts() {
  [[ "$SYNC_TREES" == "1" ]] || return 0
  [[ -n "$SYNC_HOSTS" ]] || return 0

  local src_dir="$TREE_DIR"
  local dst_dir="$TREE_DIR"
  if [[ "$TREE_DIR" != /* ]]; then
    src_dir="$ROOT_DIR/$TREE_DIR"
    dst_dir="$ROOT_DIR/$TREE_DIR"
  fi

  for host in $SYNC_HOSTS; do
    is_local_host "$host" && continue
    echo "Sync tree dir to $host:$dst_dir"
    "$SSH_BIN" "$host" "mkdir -p $(printf '%q' "$dst_dir")"
    if command -v rsync >/dev/null 2>&1; then
      rsync -a "$src_dir/" "$host:$dst_dir/"
    else
      scp -q "$src_dir"/*.pkl "$host:$dst_dir/"
    fi
  done
}

tools/manage_tn_dask_cluster.sh start

echo "Search with dask: $DASK_ADDRESS"
"$PYTHON_BIN" -u tools/tn_contest_runner.py search \
  --case "$CASE" \
  --nqubits "$NQUBITS" \
  --nlayers "$NLAYERS" \
  --observables $OBSERVABLES \
  --tree-dir "$TREE_DIR" \
  --dask-address "$DASK_ADDRESS" \
  --torch-threads "$TORCH_THREADS" \
  --dtype "$DTYPE" \
  --tn-search-repeats "$SEARCH_REPEATS" \
  --tn-search-time "$SEARCH_TIME" \
  "${tn_slice_args[@]}"

sync_trees_to_hosts

echo "Contract with MPI: $MPIEXEC_FULL"
read -r -a mpi_prefix <<< "$MPIEXEC_FULL"
"${mpi_prefix[@]}" "$PYTHON_BIN" -u tools/tn_contest_runner.py contract \
  --mpi \
  --case "$CASE" \
  --nqubits "$NQUBITS" \
  --nlayers "$NLAYERS" \
  --observables $OBSERVABLES \
  --tree-dir "$TREE_DIR" \
  --torch-threads "$TORCH_THREADS" \
  --dtype "$DTYPE" \
  "${tn_slice_args[@]}"