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

workspace="/mnt/pvc/users/simon/agentptb/runs/sol-max-s1/workspace"
eval_marker="$workspace/state/step150-eval-complete"
eval_session="agentptb-step150-eval-watcher"
train_session="agentptb-sft-continuation"
target_steps=1800
source_output="$workspace/outputs/sft-agent-mix-clean-full-v1"
continuation_log="$workspace/logs/sft-agent-mix-clean-continuation-launch.log"
full_config="$workspace/configs/sft-agent-mix-clean-full-step1650-resume.toml"
selective_config="$workspace/configs/sft-agent-mix-clean-full-step1650-selective-resume.toml"
gpu_optim_config="$workspace/configs/sft-agent-mix-clean-full-step1650-gpuoptim-resume.toml"
noact_config="$workspace/configs/sft-agent-mix-clean-full-step1650-noactoffload-resume.toml"
mb10_config="$workspace/configs/sft-agent-mix-clean-full-mb10-resume.toml"
mb12_config="$workspace/configs/sft-agent-mix-clean-full-mb12-resume.toml"
bench_pid=""

cleanup_bench() {
  if [[ -n "$bench_pid" ]]; then
    # Target the whole setsid process group even if the uv group leader has
    # already exited; torchrun/workers can otherwise survive an OOM launcher.
    kill -TERM -- "-$bench_pid" 2>/dev/null || true
    for _ in {1..12}; do
      kill -0 -- "-$bench_pid" 2>/dev/null || break
      sleep 5
    done
    kill -KILL -- "-$bench_pid" 2>/dev/null || true
    wait "$bench_pid" 2>/dev/null || true
  fi
}

run_benchmark() {
  local config="$1" output="$2" log="$3"
  mkdir -p "$output/checkpoints" "$output/weights"
  if [[ ! -e "$output/checkpoints/step_150" ]]; then
    cp -Rl "$source_output/checkpoints/step_150" "$output/checkpoints/step_150"
  fi
  if [[ ! -e "$output/weights/step_150" ]]; then
    cp -Rl "$source_output/weights/step_150" "$output/weights/step_150"
  fi
  if [[ -f "$output/checkpoints/step_152/trainer/.metadata" \
        && -f "$output/weights/step_152/STABLE" \
        && -f "$log" ]] \
    && grep -q 'Step 152 |' "$log"; then
    return 0
  fi

  cd /root/work/a/prime-rl
  setsid bash -c 'exec "$@"' _ \
    env CUDA_VISIBLE_DEVICES=4,5,6,7 \
    TMPDIR=/tmp/agentptb-s1 \
    PYTHONUNBUFFERED=1 \
    uv run --no-sync sft @ "$config" >"$log" 2>&1 &
  bench_pid=$!

  local ok=false
  for _ in {1..180}; do
    if [[ -f "$output/checkpoints/step_152/trainer/.metadata" \
          && -f "$output/weights/step_152/STABLE" ]] \
      && grep -q 'Step 152 |' "$log"; then
      ok=true
      break
    fi
    kill -0 "$bench_pid" 2>/dev/null || break
    sleep 10
  done
  cleanup_bench
  bench_pid=""
  [[ "$ok" == true ]]
}

read_metrics() {
  python - "$1" "$2" "$3" <<'PY'
import json
import math
import re
import sys

text = open(sys.argv[1], errors="replace").read()
metrics_path = sys.argv[2]
batch_size = int(sys.argv[3])
matches = re.findall(
    r"Step 152 \|.*?Loss ([^ |]+) \| Grad\. Norm ([^ |]+) \| "
    r"LR ([^ |]+) \| Throughput [0-9.]+ tokens/s.*?Peak Mem\. ([0-9.]+)/",
    text,
)
if not matches:
    print("0 999 nan nan nan 0")
else:
    loss, grad, lr, memory = matches[-1]
    step_metrics = {}
    try:
        with open(metrics_path) as handle:
            for line in handle:
                row = json.loads(line)
                if row.get("step") == 152:
                    step_metrics.update(row)
        effective_step_seconds = (
            float(step_metrics["time/step"])
            - float(step_metrics.get("time/save_ckpt", 0.0))
        )
        throughput = batch_size * 32768 / effective_step_seconds
    except (OSError, ValueError, KeyError, ZeroDivisionError, json.JSONDecodeError):
        effective_step_seconds = math.nan
        throughput = 0.0
    values = [
        float(loss),
        float(grad),
        float(lr),
        float(throughput),
        float(memory),
        effective_step_seconds,
    ]
    valid = int(
        all(math.isfinite(value) for value in values)
        and 1.9e-5 <= values[2] <= 2.1e-5
        and values[4] <= 170
        and values[5] > 0
    )
    print(f"{throughput:.3f}", memory, loss, grad, lr, valid)
PY
}
trap cleanup_bench EXIT INT TERM

while [[ ! -f "$eval_marker" ]]; do
  sleep 30
done
while tmux has-session -t "$eval_session" 2>/dev/null; do
  sleep 10
done

[[ -f "$source_output/checkpoints/step_150/trainer/.metadata" ]]
[[ -f "$source_output/weights/step_150/STABLE" ]]
if curl --fail --silent http://127.0.0.1:8200/health >/dev/null 2>&1; then
  echo "refusing to train while inference still owns port 8200" >&2
  exit 1
fi
tmux has-session -t "$train_session" 2>/dev/null && {
  echo "continuation session already exists" >&2
  exit 1
}

baseline_save=false
baseline_throughput=0
baseline_memory=999
baseline_loss=nan
baseline_grad=nan
baseline_lr=nan
baseline_valid=0
if run_benchmark \
  "$workspace/configs/sft-baseline-bench-resume150.toml" \
  "$workspace/outputs/baseline-bench" \
  "$workspace/logs/sft-baseline-bench.log"; then
  baseline_save=true
  read -r baseline_throughput baseline_memory baseline_loss baseline_grad baseline_lr baseline_valid \
    < <(read_metrics "$workspace/logs/sft-baseline-bench.log" \
      "$workspace/outputs/baseline-bench/metrics.jsonl" 128)
fi

mb10_save=false
mb10_throughput=0
mb10_memory=999
mb10_loss=nan
mb10_grad=nan
mb10_lr=nan
mb10_valid=0
if run_benchmark \
  "$workspace/configs/sft-mb10-bench-resume150.toml" \
  "$workspace/outputs/mb10-bench" \
  "$workspace/logs/sft-mb10-bench.log"; then
  mb10_save=true
  read -r mb10_throughput mb10_memory mb10_loss mb10_grad mb10_lr mb10_valid \
    < <(read_metrics "$workspace/logs/sft-mb10-bench.log" \
      "$workspace/outputs/mb10-bench/metrics.jsonl" 120)
fi

mb12_save=false
mb12_throughput=0
mb12_memory=999
mb12_loss=nan
mb12_grad=nan
mb12_lr=nan
mb12_valid=0
if run_benchmark \
  "$workspace/configs/sft-mb12-bench-resume150.toml" \
  "$workspace/outputs/mb12-bench" \
  "$workspace/logs/sft-mb12-bench.log"; then
  mb12_save=true
  read -r mb12_throughput mb12_memory mb12_loss mb12_grad mb12_lr mb12_valid \
    < <(read_metrics "$workspace/logs/sft-mb12-bench.log" \
      "$workspace/outputs/mb12-bench/metrics.jsonl" 96)
fi

selective_save=false
selective_throughput=0
selective_memory=999
selective_loss=nan
selective_grad=nan
selective_lr=nan
selective_valid=0
if run_benchmark \
  "$workspace/configs/sft-ac-selective-bench-resume150.toml" \
  "$workspace/outputs/ac-selective-bench" \
  "$workspace/logs/sft-ac-selective-bench.log"; then
  selective_save=true
  read -r selective_throughput selective_memory selective_loss selective_grad selective_lr selective_valid \
    < <(read_metrics "$workspace/logs/sft-ac-selective-bench.log" \
      "$workspace/outputs/ac-selective-bench/metrics.jsonl" 128)
fi

gpu_save=false
gpu_throughput=0
gpu_memory=999
gpu_loss=nan
gpu_grad=nan
gpu_lr=nan
gpu_valid=0
if run_benchmark \
  "$workspace/configs/sft-optim-gpu-bench-resume150.toml" \
  "$workspace/outputs/optim-gpu-bench" \
  "$workspace/logs/sft-optim-gpu-bench.log"; then
  gpu_save=true
  read -r gpu_throughput gpu_memory gpu_loss gpu_grad gpu_lr gpu_valid \
    < <(read_metrics "$workspace/logs/sft-optim-gpu-bench.log" \
      "$workspace/outputs/optim-gpu-bench/metrics.jsonl" 128)
fi

noact_save=false
noact_throughput=0
noact_memory=999
noact_loss=nan
noact_grad=nan
noact_lr=nan
noact_valid=0
if run_benchmark \
  "$workspace/configs/sft-noactoffload-bench-resume150.toml" \
  "$workspace/outputs/noactoffload-bench" \
  "$workspace/logs/sft-noactoffload-bench.log"; then
  noact_save=true
  read -r noact_throughput noact_memory noact_loss noact_grad noact_lr noact_valid \
    < <(read_metrics "$workspace/logs/sft-noactoffload-bench.log" \
      "$workspace/outputs/noactoffload-bench/metrics.jsonl" 128)
fi

chosen="$full_config"
best_throughput="$baseline_throughput"
chosen_batch_size=128
mb10_eligible=0
mb12_eligible=0
selective_eligible=0
gpu_eligible=0
noact_eligible=0
if [[ "$baseline_valid" == 1 && "$mb10_valid" == 1 ]] \
  && awk -v new="$mb10_throughput" -v base="$baseline_throughput" \
    'BEGIN { exit !(new >= base * 1.03) }'; then
  mb10_eligible=1
  chosen="$mb10_config"
  chosen_batch_size=120
  best_throughput="$mb10_throughput"
fi
if [[ "$baseline_valid" == 1 && "$mb12_valid" == 1 ]] \
  && awk -v new="$mb12_throughput" -v base="$baseline_throughput" \
    'BEGIN { exit !(new >= base * 1.03) }'; then
  mb12_eligible=1
  if awk -v new="$mb12_throughput" -v old="$best_throughput" \
    'BEGIN { exit !(new > old) }'; then
    chosen="$mb12_config"
    chosen_batch_size=96
    best_throughput="$mb12_throughput"
  fi
fi
if [[ "$baseline_valid" == 1 && "$selective_valid" == 1 ]] \
  && awk -v new="$selective_throughput" -v base="$baseline_throughput" \
    'BEGIN { exit !(new >= base * 1.03) }'; then
  selective_eligible=1
  if awk -v new="$selective_throughput" -v old="$best_throughput" \
    'BEGIN { exit !(new > old) }'; then
    chosen="$selective_config"
    chosen_batch_size=128
    best_throughput="$selective_throughput"
  fi
fi
if [[ "$baseline_valid" == 1 && "$gpu_valid" == 1 ]] \
  && awk -v new="$gpu_throughput" -v base="$baseline_throughput" \
    'BEGIN { exit !(new >= base * 1.03) }'; then
  gpu_eligible=1
  if awk -v new="$gpu_throughput" -v old="$best_throughput" \
    'BEGIN { exit !(new > old) }'; then
    chosen="$gpu_optim_config"
    chosen_batch_size=128
    best_throughput="$gpu_throughput"
  fi
fi
if [[ "$baseline_valid" == 1 && "$noact_valid" == 1 ]] \
  && awk -v new="$noact_throughput" -v base="$baseline_throughput" \
    'BEGIN { exit !(new >= base * 1.03) }'; then
  noact_eligible=1
  if awk -v new="$noact_throughput" -v old="$best_throughput" \
    'BEGIN { exit !(new > old) }'; then
    chosen="$noact_config"
    chosen_batch_size=128
    best_throughput="$noact_throughput"
  fi
fi
trap - EXIT INT TERM

# Choose the terminal step only from measured execution throughput and remaining
# wall clock.  This decision is deliberately independent of evaluation reward.
# Reserve 18 hours for checkpoint scanning, serving, the canonical 1,178-episode
# stock/custom measurement, and recovery from transient infrastructure failures.
# A 5% step-time penalty plus 1.5 seconds/update for amortized checkpoint overhead
# keeps a two-update benchmark from making the long-run projection optimistic.
remaining_seconds="$(python - "$workspace/DEADLINE" <<'PY'
import sys
import time

print(max(0, int(float(open(sys.argv[1]).read().strip()) - time.time())))
PY
)"
planning_throughput="$best_throughput"
if ! awk -v value="$planning_throughput" \
  'BEGIN { exit !(value ~ /^[0-9]+([.][0-9]+)?$/ && value >= 1000) }'; then
  # The active run sustained roughly 27.5k tokens/s.  Use a lower fallback if
  # the isolated reference benchmark failed before emitting a usable metric.
  planning_throughput=26000
fi
final_eval_reserve_seconds=$((18 * 3600))
startup_guard_seconds=1800
read -r target_steps projected_step_seconds < <(
  python - "$remaining_seconds" "$planning_throughput" \
    "$final_eval_reserve_seconds" "$startup_guard_seconds" \
    "$chosen_batch_size" <<'PY'
import sys

remaining = int(sys.argv[1])
throughput = float(sys.argv[2])
reserve = int(sys.argv[3])
startup_guard = int(sys.argv[4])
batch_size = int(sys.argv[5])
tokens_per_update = batch_size * 32768
projected_step_seconds = tokens_per_update / throughput * 1.05 + 1.5
available = max(0, remaining - reserve - startup_guard)
updates = int(available / projected_step_seconds)
# Step 150 is the fixed branch point.  About step 2,050 is two full passes over
# the prepared mixture under batch 128. Preserve the same total-token cap when
# a mechanically faster benchmark selects a nearby global batch size.
token_equivalent_cap = 150 + int((2050 - 150) * 128 / batch_size)
target = max(151, min(token_equivalent_cap, 150 + updates))
print(target, f"{projected_step_seconds:.3f}")
PY
)

printf 'baseline_checkpoint_save_succeeded=%s\nbaseline_throughput=%s\nbaseline_peak_memory_gib=%s\nbaseline_loss=%s\nbaseline_grad_norm=%s\nbaseline_lr=%s\nbaseline_valid=%s\nmb10_checkpoint_save_succeeded=%s\nmb10_throughput=%s\nmb10_peak_memory_gib=%s\nmb10_loss=%s\nmb10_grad_norm=%s\nmb10_lr=%s\nmb10_valid=%s\nmb10_eligible=%s\nmb12_checkpoint_save_succeeded=%s\nmb12_throughput=%s\nmb12_peak_memory_gib=%s\nmb12_loss=%s\nmb12_grad_norm=%s\nmb12_lr=%s\nmb12_valid=%s\nmb12_eligible=%s\nselective_checkpoint_save_succeeded=%s\nselective_throughput=%s\nselective_peak_memory_gib=%s\nselective_loss=%s\nselective_grad_norm=%s\nselective_lr=%s\nselective_valid=%s\nselective_eligible=%s\ngpu_optim_checkpoint_save_succeeded=%s\ngpu_optim_throughput=%s\ngpu_optim_peak_memory_gib=%s\ngpu_optim_loss=%s\ngpu_optim_grad_norm=%s\ngpu_optim_lr=%s\ngpu_optim_valid=%s\ngpu_optim_eligible=%s\nnoact_checkpoint_save_succeeded=%s\nnoact_throughput=%s\nnoact_peak_memory_gib=%s\nnoact_loss=%s\nnoact_grad_norm=%s\nnoact_lr=%s\nnoact_valid=%s\nnoact_eligible=%s\nchosen_config=%s\nchosen_batch_size=%s\nremaining_seconds_at_plan=%s\nplanning_throughput=%s\nprojected_step_seconds=%s\nfinal_eval_reserve_seconds=%s\nstartup_guard_seconds=%s\ntarget_steps=%s\n' \
  "$baseline_save" "$baseline_throughput" "$baseline_memory" "$baseline_loss" "$baseline_grad" "$baseline_lr" "$baseline_valid" \
  "$mb10_save" "$mb10_throughput" "$mb10_memory" "$mb10_loss" "$mb10_grad" "$mb10_lr" "$mb10_valid" "$mb10_eligible" \
  "$mb12_save" "$mb12_throughput" "$mb12_memory" "$mb12_loss" "$mb12_grad" "$mb12_lr" "$mb12_valid" "$mb12_eligible" \
  "$selective_save" "$selective_throughput" "$selective_memory" "$selective_loss" "$selective_grad" "$selective_lr" "$selective_valid" "$selective_eligible" \
  "$gpu_save" "$gpu_throughput" "$gpu_memory" "$gpu_loss" "$gpu_grad" "$gpu_lr" "$gpu_valid" "$gpu_eligible" \
  "$noact_save" "$noact_throughput" "$noact_memory" "$noact_loss" "$noact_grad" "$noact_lr" "$noact_valid" "$noact_eligible" "$chosen" "$chosen_batch_size" \
  "$remaining_seconds" "$planning_throughput" "$projected_step_seconds" \
  "$final_eval_reserve_seconds" "$startup_guard_seconds" "$target_steps" \
  >"$workspace/state/training-variant-benchmarks.txt"

tmux has-session -t "$train_session" 2>/dev/null && {
  echo "continuation session already exists" >&2
  exit 1
}
tmux new-session -d -s "$train_session" \
  "cd /root/work/a/prime-rl && env CUDA_VISIBLE_DEVICES=4,5,6,7 TMPDIR=/tmp/agentptb-s1 PYTHONUNBUFFERED=1 uv run --no-sync sft @ '$chosen' --max-steps '$target_steps' > '$continuation_log' 2>&1"
printf 'config=%s\nbatch_size=%s\nmax_steps=%s\n' "$chosen" "$chosen_batch_size" "$target_steps" \
  >"$workspace/state/continuation-config.txt"
touch "$workspace/state/continuation-launched"