#!/usr/bin/env bash
# After both P1 selectors are complete, bind them into frozen R2/T2
# worktrees, run the full 27+27 cells sequentially, audit every record, and
# commit the generated evidence on the corresponding run branches.
set -euo pipefail

if [ "$#" -ne 3 ]; then
  echo "usage: $0 EXPECTED_P1_COMMIT EXPECTED_R2_BASE EXPECTED_T2_BASE" >&2
  exit 2
fi

expected_p1_commit="$1"
expected_r2_base="$2"
expected_t2_base="$3"
resnet_p1_root="/scratch/yurenh2/sdil-runs/resnet-p1"
transformer_p1_root="/scratch/yurenh2/sdil-runs/transformer-p1"
resnet_r2_root="/scratch/yurenh2/sdil-runs/resnet-r2"
transformer_t2_root="/scratch/yurenh2/sdil-runs/transformer-t2"
plain_root="/scratch/yurenh2/sdil-crossover/dualprop_vgg16"
python_bin="/home/yurenh2/miniconda3/envs/ep_pascal/bin/python3"
plain_python="/scratch/yurenh2/venvs/dualprop/bin/python"
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"

require_source() {
  local root="$1"
  local expected_commit="$2"
  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
}

wait_for_file() {
  while [ ! -f "$1" ]; do
    sleep 60
  done
}

validate_selector() {
  local path="$1"
  local kind="$2"
  "$python_bin" "$validator" --path "$path" --kind "$kind"
}

test -x "$python_bin"
test -x "$plain_python"
test "$(gpu_uuid 5)" = "$gpu5_uuid"
test "$(gpu_uuid 7)" = "$gpu7_uuid"
require_source "$resnet_p1_root" "$expected_p1_commit"
require_source "$transformer_p1_root" "$expected_p1_commit"
require_source "$resnet_r2_root" "$expected_r2_base"
require_source "$transformer_t2_root" "$expected_t2_base"

resnet_p1_selector="$resnet_p1_root/results/resnet_crossover/p1_selector.json"
transformer_p1_selector="$transformer_p1_root/results/transformer_crossover/p1_selector.json"
echo "waiting for both audited P1 selectors at $(date --iso-8601=seconds)"
wait_for_file "$resnet_p1_selector"
wait_for_file "$transformer_p1_selector"
validate_selector "$resnet_p1_selector" "resnet_p1_selector"
validate_selector \
  "$transformer_p1_selector" "transformer_p1_selector"

resnet_r2_selector="$resnet_r2_root/results/resnet_crossover/p1_selector.json"
transformer_t2_selector="$transformer_t2_root/results/transformer_crossover/p1_selector.json"
test ! -e "$resnet_r2_selector"
test ! -e "$transformer_t2_selector"
mkdir -p "$(dirname "$resnet_r2_selector")"
mkdir -p "$(dirname "$transformer_t2_selector")"
(
  cd "$resnet_p1_root"
  "$python_bin" experiments/analyze_resnet_crossover_p1.py \
    --out "$resnet_r2_selector"
) >"$resnet_r2_root/results/resnet_crossover/p1_import.log" 2>&1
(
  cd "$transformer_p1_root"
  "$python_bin" experiments/analyze_transformer_crossover_p1.py \
    --out "$transformer_t2_selector"
) >"$transformer_t2_root/results/transformer_crossover/p1_import.log" 2>&1
validate_selector "$resnet_r2_selector" "resnet_p1_selector"
validate_selector \
  "$transformer_t2_selector" "transformer_p1_selector"
git -C "$resnet_r2_root" add -f \
  results/resnet_crossover/p1_selector.json
git -C "$resnet_r2_root" commit -m \
  "experiment: freeze ResNet crossover selector"
git -C "$transformer_t2_root" add -f \
  results/transformer_crossover/p1_selector.json
git -C "$transformer_t2_root" commit -m \
  "experiment: freeze Transformer crossover selector"
resnet_r2_commit="$(git -C "$resnet_r2_root" rev-parse HEAD)"
transformer_t2_commit="$(git -C "$transformer_t2_root" rev-parse HEAD)"
require_source "$resnet_r2_root" "$resnet_r2_commit"
require_source "$transformer_t2_root" "$transformer_t2_commit"
echo "selectors frozen; launching ResNet R2 at $(date --iso-8601=seconds)"

wait_for_idle_gpus
if tmux has-session -t resnet_r2_gpu5 2>/dev/null ||
   tmux has-session -t resnet_r2_gpu7 2>/dev/null; then
  echo "ResNet R2 session already exists; refusing duplicate launch" >&2
  exit 1
fi
tmux new-session -d -s resnet_r2_gpu5 \
  "cd '$resnet_r2_root' && env CUDA_VISIBLE_DEVICES=5 '$python_bin' experiments/resnet_crossover_r2.py --num-shards 2 --shard-index 0 > results/resnet_crossover/r2_gpu5.log 2>&1"
resnet_r2_launch="$resnet_r2_root/results/resnet_crossover/r2_launch.json"
until "$python_bin" "$validator" \
    --path "$resnet_r2_launch" \
    --kind resnet_r2_launch \
    --expected-commit "$resnet_r2_commit" >/dev/null 2>&1; do
  sleep 1
done
tmux new-session -d -s resnet_r2_gpu7 \
  "cd '$resnet_r2_root' && env CUDA_VISIBLE_DEVICES=7 '$python_bin' experiments/resnet_crossover_r2.py --num-shards 2 --shard-index 1 > results/resnet_crossover/r2_gpu7.log 2>&1"
tmux has-session -t resnet_r2_gpu5
tmux has-session -t resnet_r2_gpu7
while tmux has-session -t resnet_r2_gpu5 2>/dev/null ||
      tmux has-session -t resnet_r2_gpu7 2>/dev/null; do
  sleep 60
done
test "$(manifest_count \
  "$resnet_r2_root/results/resnet_crossover/r2")" -eq 27
require_source "$resnet_r2_root" "$resnet_r2_commit"
(
  cd "$resnet_r2_root"
  "$python_bin" experiments/analyze_resnet_crossover_r2.py \
    --out results/resnet_crossover/r2_audit.json
) >"$resnet_r2_root/results/resnet_crossover/r2_audit.log" 2>&1
"$python_bin" "$validator" \
  --path "$resnet_r2_root/results/resnet_crossover/r2_audit.json" \
  --kind resnet_r2_audit
git -C "$resnet_r2_root" add -f results/resnet_crossover
git -C "$resnet_r2_root" commit -m \
  "results: record complete ResNet crossover"
echo "ResNet R2 audited and committed at $(date --iso-8601=seconds)"

wait_for_idle_gpus
require_source "$transformer_t2_root" "$transformer_t2_commit"
if tmux has-session -t transformer_t2_gpu5 2>/dev/null ||
   tmux has-session -t transformer_t2_gpu7 2>/dev/null; then
  echo "Transformer T2 session already exists; refusing duplicate launch" >&2
  exit 1
fi
tmux new-session -d -s transformer_t2_gpu5 \
  "cd '$transformer_t2_root' && env CUDA_VISIBLE_DEVICES=5 '$python_bin' experiments/transformer_crossover_t2.py --num-shards 2 --shard-index 0 > results/transformer_crossover/t2_gpu5.log 2>&1"
transformer_t2_launch="$transformer_t2_root/results/transformer_crossover/t2_launch.json"
until "$python_bin" "$validator" \
    --path "$transformer_t2_launch" \
    --kind transformer_t2_launch \
    --expected-commit "$transformer_t2_commit" >/dev/null 2>&1; do
  sleep 1
done
tmux new-session -d -s transformer_t2_gpu7 \
  "cd '$transformer_t2_root' && env CUDA_VISIBLE_DEVICES=7 '$python_bin' experiments/transformer_crossover_t2.py --num-shards 2 --shard-index 1 > results/transformer_crossover/t2_gpu7.log 2>&1"
tmux has-session -t transformer_t2_gpu5
tmux has-session -t transformer_t2_gpu7
while tmux has-session -t transformer_t2_gpu5 2>/dev/null ||
      tmux has-session -t transformer_t2_gpu7 2>/dev/null; do
  sleep 60
done
test "$(manifest_count \
  "$transformer_t2_root/results/transformer_crossover/t2")" -eq 27
require_source "$transformer_t2_root" "$transformer_t2_commit"
(
  cd "$transformer_t2_root"
  "$python_bin" experiments/analyze_transformer_crossover_t2.py \
    --out results/transformer_crossover/t2_audit.json
) >"$transformer_t2_root/results/transformer_crossover/t2_audit.log" 2>&1
"$python_bin" "$validator" \
  --path "$transformer_t2_root/results/transformer_crossover/t2_audit.json" \
  --kind transformer_t2_audit

plain_audit="$transformer_t2_root/results/plain_cnn_p2_audit.json"
(
  cd "$plain_root"
  "$plain_python" analyze_p2.py --output "$plain_audit"
)
"$python_bin" "$transformer_t2_root/experiments/analyze_crossover_81.py" \
  --plain "$plain_audit" \
  --resnet "$resnet_r2_root/results/resnet_crossover/r2_audit.json" \
  --transformer \
    "$transformer_t2_root/results/transformer_crossover/t2_audit.json" \
  --out "$transformer_t2_root/results/crossover_81_audit.json"
"$python_bin" "$validator" \
  --path "$transformer_t2_root/results/crossover_81_audit.json" \
  --kind crossover_81_audit
git -C "$transformer_t2_root" add -f \
  results/transformer_crossover \
  results/plain_cnn_p2_audit.json \
  results/crossover_81_audit.json
git -C "$transformer_t2_root" commit -m \
  "results: record complete 81-cell crossover"
echo "complete 81-cell crossover audited and committed at $(date --iso-8601=seconds)"