#!/usr/bin/env bash
# Wait for a complete frozen ResNet P1, audit its selector, then launch and
# audit Transformer P1 on the same two authorized physical GPUs. The watcher
# never stops a process and refuses incomplete shards, source drift, or a
# duplicate launch.
set -euo pipefail

if [ "$#" -ne 2 ]; then
  echo "usage: $0 EXPECTED_RESNET_COMMIT EXPECTED_TRANSFORMER_COMMIT" >&2
  exit 2
fi

expected_resnet_commit="$1"
expected_transformer_commit="$2"
resnet_root="/scratch/yurenh2/sdil-runs/resnet-p1"
transformer_root="/scratch/yurenh2/sdil-runs/transformer-p1"
python_bin="/home/yurenh2/miniconda3/envs/ep_pascal/bin/python3"
validator="/home/yurenh2/sdil/experiments/validate_crossover_artifact.py"
gpu5_uuid="GPU-f4aed756-228e-7522-94e7-fb543c02cc74"
gpu7_uuid="GPU-df3d5933-c3e9-9ea1-d595-9d2a3b68c14b"
resnet_launch="$resnet_root/results/resnet_crossover/p1_launch.json"
transformer_launch="$transformer_root/results/transformer_crossover/p1_launch.json"

require_source() {
  local root="$1"
  local expected_commit="$2"
  test -d "$root/.git" -o -f "$root/.git"
  test "$(git -C "$root" rev-parse HEAD)" = "$expected_commit"
  test -z "$(git -C "$root" status --porcelain --untracked-files=no)"
}

gpu_uuid() {
  nvidia-smi \
    --query-gpu=index,uuid \
    --format=csv,noheader,nounits |
    awk -F',' -v wanted="$1" '
      {
        gsub(/^[ \t]+|[ \t]+$/, "", $1)
        gsub(/^[ \t]+|[ \t]+$/, "", $2)
        if ($1 == wanted) print $2
      }'
}

gpu_has_compute_process() {
  nvidia-smi \
    --query-compute-apps=gpu_uuid \
    --format=csv,noheader,nounits 2>/dev/null |
    awk -v wanted="$1" '
      {
        gsub(/^[ \t]+|[ \t]+$/, "", $0)
        if ($0 == wanted) found=1
      }
      END { exit(found ? 0 : 1) }'
}

wait_for_idle_gpus() {
  while gpu_has_compute_process "$gpu5_uuid" ||
        gpu_has_compute_process "$gpu7_uuid"; do
    sleep 60
  done
}

manifest_count() {
  local result_root="$1"
  find "$result_root" -maxdepth 1 -type f \
    -name '*.manifest.json' -print 2>/dev/null | wc -l
}

require_source "$resnet_root" "$expected_resnet_commit"
require_source "$transformer_root" "$expected_transformer_commit"
test -x "$python_bin"
test "$(gpu_uuid 5)" = "$gpu5_uuid"
test "$(gpu_uuid 7)" = "$gpu7_uuid"

echo "waiting for frozen ResNet P1 launch at $(date --iso-8601=seconds)"
while [ ! -f "$resnet_launch" ]; do
  sleep 60
done
"$python_bin" "$validator" \
  --path "$resnet_launch" \
  --kind resnet_p1_launch \
  --expected-commit "$expected_resnet_commit"

echo "ResNet P1 launch observed; waiting for both shards"
while tmux has-session -t resnet_p1_gpu5 2>/dev/null ||
      tmux has-session -t resnet_p1_gpu7 2>/dev/null; do
  sleep 60
done
test "$(manifest_count \
  "$resnet_root/results/resnet_crossover/p1")" -eq 19
require_source "$resnet_root" "$expected_resnet_commit"
(
  cd "$resnet_root"
  "$python_bin" experiments/analyze_resnet_crossover_p1.py \
    --out results/resnet_crossover/p1_selector.json
) >"$resnet_root/results/resnet_crossover/p1_selector.log" 2>&1
"$python_bin" "$validator" \
  --path "$resnet_root/results/resnet_crossover/p1_selector.json" \
  --kind resnet_p1_selector
echo "ResNet P1 selector audited at $(date --iso-8601=seconds)"

wait_for_idle_gpus
require_source "$transformer_root" "$expected_transformer_commit"
test "$(gpu_uuid 5)" = "$gpu5_uuid"
test "$(gpu_uuid 7)" = "$gpu7_uuid"
if tmux has-session -t transformer_p1_gpu5 2>/dev/null ||
   tmux has-session -t transformer_p1_gpu7 2>/dev/null ||
   [ -e "$transformer_launch" ]; then
  echo "Transformer P1 already launched; refusing duplicate launch" >&2
  exit 1
fi

mkdir -p "$transformer_root/results/transformer_crossover"
tmux new-session -d -s transformer_p1_gpu5 \
  "cd '$transformer_root' && env CUDA_VISIBLE_DEVICES=5 '$python_bin' experiments/transformer_crossover_grid.py --stage p1 --num-shards 2 --shard-index 0 > results/transformer_crossover/p1_gpu5.log 2>&1"
until "$python_bin" "$validator" \
    --path "$transformer_launch" \
    --kind transformer_p1_launch \
    --expected-commit "$expected_transformer_commit" >/dev/null 2>&1; do
  sleep 1
done
tmux new-session -d -s transformer_p1_gpu7 \
  "cd '$transformer_root' && env CUDA_VISIBLE_DEVICES=7 '$python_bin' experiments/transformer_crossover_grid.py --stage p1 --num-shards 2 --shard-index 1 > results/transformer_crossover/p1_gpu7.log 2>&1"
tmux has-session -t transformer_p1_gpu5
tmux has-session -t transformer_p1_gpu7
echo "launched frozen Transformer P1 shards at $(date --iso-8601=seconds)"

while tmux has-session -t transformer_p1_gpu5 2>/dev/null ||
      tmux has-session -t transformer_p1_gpu7 2>/dev/null; do
  sleep 60
done
test "$(manifest_count \
  "$transformer_root/results/transformer_crossover/p1")" -eq 27
require_source "$transformer_root" "$expected_transformer_commit"
(
  cd "$transformer_root"
  "$python_bin" experiments/analyze_transformer_crossover_p1.py \
    --out results/transformer_crossover/p1_selector.json
) >"$transformer_root/results/transformer_crossover/p1_selector.log" 2>&1
"$python_bin" "$validator" \
  --path "$transformer_root/results/transformer_crossover/p1_selector.json" \
  --kind transformer_p1_selector
echo "Transformer P1 selector audited at $(date --iso-8601=seconds)"