#!/usr/bin/env bash
# Mini per-position curves chain (uli decision 2026-08-06, option A):
# re-runs the 20 CONFIG-WITH-DRAFTER configs (12 Qwen + 8 Gemma) on the
# stratified 180-prompt subset (acc-sample.jsonl) with FRESH dirs
# experiments/runs/curves-<name>/ to recover "acc per pos" (server.log).
#
# Why this is valid: T=0 + seed 42 + the 180 prompts are a subset of the
# 1474 (f1-sample) → generations IDENTICAL to the full run → the subset
# curve IS the curve of those prompts in the real run (deterministic
# equivalence; documented in the Methods of the post/paper).
#
# Target-solo configs are NOT included (no drafter → no curves).
# The 8 Gemma-draft configs ALREADY have complete curves (n=1474) in their
# server.log; this mini-chain re-runs them at n=180 for a uniform F1 figure
# (same n in every panel).
#
# Usage:
#   DRY_RUN=1 bash scripts/run_curves_chain.sh   # prints the matrix (20) and exits
#   bash scripts/run_curves_chain.sh             # ~15 min/config ≈ 4.5 h total
#
# RUN ONLY WITH THE GPU FREE (after "FINAL RUN DONE" of the full chain).
set -uo pipefail
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$REPO"
source scripts/spec-env.sh
set +e  # spec-env.sh enables set -e; the chain handles errors via exit codes

LOG="experiments/runs/queue-curves.log"
log() { echo "[$(date '+%F %T')] $*" | tee -a "$LOG"; }

SAMPLES="experiments/prompts/old/acc-sample.jsonl"
NTOK=256
Q="models/Qwen3-8B"
G="models/gemma-4-12b-it"
DK="models/drafts"
BENCH="python scripts/bench_accept.py"

OK_COUNT=0
FAILED_COUNT=0
SKIP_COUNT=0
OOM_OK=1
G8_CHECKED=0

# Matrix of 20 configs with drafter (spec final-run-matrix): name|target|draft|spec_type|p_min|out|extra
MATRIX=(
    # Qwen3-8B (12)
    "curves-qwen-q4-vanilla17b|$Q-Q4_K_M.gguf|$DK/Qwen3-1.7B-Q4_K_M.gguf|draft-simple||experiments/runs/curves-qwen-q4-vanilla17b"
    "curves-qwen-q4-eagle3|$Q-Q4_K_M.gguf|$DK/Qwen3-8B-speculator.eagle3-F16.gguf|draft-eagle3||experiments/runs/curves-qwen-q4-eagle3"
    "curves-qwen-q4-dflash|$Q-Q4_K_M.gguf|$DK/dflash_qwen3_8b_block7.gguf|draft-dflash||experiments/runs/curves-qwen-q4-dflash"
    "curves-qwen-q4-dspark-p0|$Q-Q4_K_M.gguf|$DK/dspark_qwen3_8b_block7.gguf|draft-dspark|0.0|experiments/runs/curves-qwen-q4-dspark-p0"
    "curves-qwen-q4-dspark-p2|$Q-Q4_K_M.gguf|$DK/dspark_qwen3_8b_block7.gguf|draft-dspark|0.2|experiments/runs/curves-qwen-q4-dspark-p2"
    "curves-qwen-q4-dspark-p4|$Q-Q4_K_M.gguf|$DK/dspark_qwen3_8b_block7.gguf|draft-dspark|0.4|experiments/runs/curves-qwen-q4-dspark-p4"
    "curves-qwen-q4-dspark-p6|$Q-Q4_K_M.gguf|$DK/dspark_qwen3_8b_block7.gguf|draft-dspark|0.6|experiments/runs/curves-qwen-q4-dspark-p6"
    "curves-qwen-q5-eagle3|$Q-Q5_K_M.gguf|$DK/Qwen3-8B-speculator.eagle3-F16.gguf|draft-eagle3||experiments/runs/curves-qwen-q5-eagle3"
    "curves-qwen-q5-dflash|$Q-Q5_K_M.gguf|$DK/dflash_qwen3_8b_block7.gguf|draft-dflash||experiments/runs/curves-qwen-q5-dflash"
    "curves-qwen-q5-dspark-p0|$Q-Q5_K_M.gguf|$DK/dspark_qwen3_8b_block7.gguf|draft-dspark|0.0|experiments/runs/curves-qwen-q5-dspark-p0"
    "curves-qwen-q8-eagle3|$Q-Q8_0.gguf|$DK/Qwen3-8B-speculator.eagle3-F16.gguf|draft-eagle3||experiments/runs/curves-qwen-q8-eagle3"
    "curves-qwen-q8-dspark-p0|$Q-Q8_0.gguf|$DK/dspark_qwen3_8b_block7.gguf|draft-dspark|0.0|experiments/runs/curves-qwen-q8-dspark-p0"
    # Gemma 4 12B (8)
    "curves-gemma-q4-mtp|$G-Q4_K_M.gguf|$DK/mtp-gemma-4-12b-it-Q8_0.gguf|draft-mtp||experiments/runs/curves-gemma-q4-mtp"
    "curves-gemma-q4-dflash-f16|$G-Q4_K_M.gguf|$DK/gemma-4-12B-it-DFlash-F16.gguf|draft-dflash||experiments/runs/curves-gemma-q4-dflash-f16"
    "curves-gemma-q4-dflash-q4|$G-Q4_K_M.gguf|$DK/gemma-4-12B-it-DFlash-Q4_K_M.gguf|draft-dflash||experiments/runs/curves-gemma-q4-dflash-q4"
    "curves-gemma-q4-dflash-q8|$G-Q4_K_M.gguf|$DK/gemma-4-12B-it-DFlash-Q8_0.gguf|draft-dflash||experiments/runs/curves-gemma-q4-dflash-q8"
    "curves-gemma-q5-mtp|$G-Q5_K_M.gguf|$DK/mtp-gemma-4-12b-it-Q8_0.gguf|draft-mtp||experiments/runs/curves-gemma-q5-mtp"
    "curves-gemma-q5-dflash-f16|$G-Q5_K_M.gguf|$DK/gemma-4-12B-it-DFlash-F16.gguf|draft-dflash||experiments/runs/curves-gemma-q5-dflash-f16"
    "curves-gemma-q8-mtp|$G-Q8_0.gguf|$DK/mtp-gemma-4-12b-it-Q8_0.gguf|draft-mtp||experiments/runs/curves-gemma-q8-mtp"
    "curves-gemma-q8-dflash-f16|$G-Q8_0.gguf|$DK/gemma-4-12B-it-DFlash-F16.gguf|draft-dflash||experiments/runs/curves-gemma-q8-dflash-f16|-ub 512"
)

have_file() {
    local f="$1"
    [ -f "$f" ] && [ "$(stat -c%s "$f" 2>/dev/null || echo 0)" -gt 100000000 ]
}

run_job() {
    local name="$1" out="$2"
    shift 2
    local rc
    log "CURVE START: $name"
    $BENCH --config-name "$name" --prompts "$SAMPLES" --n-tokens "$NTOK" \
        --out "$out" --resume "$@"
    rc=$?
    case "$rc" in
        0) log "CURVE DONE: $name (exit 0)"; OK_COUNT=$((OK_COUNT + 1)) ;;
        2) log "CURVE WITH FAILURES: $name (exit 2) — prompts in errors.jsonl"; FAILED_COUNT=$((FAILED_COUNT + 1)) ;;
        3) log "CURVE LOCK: $name (exit 3) — another runner on this --out; skipping"; FAILED_COUNT=$((FAILED_COUNT + 1)) ;;
        *) log "CURVE FAILED: $name (exit $rc)"; FAILED_COUNT=$((FAILED_COUNT + 1)) ;;
    esac
    return "$rc"
}

run_config() {
    local entry="$1"
    local name target draft stype pmin out extra
    IFS='|' read -r name target draft stype pmin out extra <<< "$entry"
    [ -n "$out" ] || out="experiments/runs/$name"
    if ! have_file "$target" || ! have_file "$draft"; then
        log "SKIP: $name (missing target/draft)"
        SKIP_COUNT=$((SKIP_COUNT + 1))
        return
    fi
    local args=(--model "$target" --draft "$draft" --spec-type "$stype")
    [ -n "$pmin" ] && args+=(--spec-draft-p-min "$pmin")
    [ -n "$extra" ] && args+=(--extra "$extra")
    run_job "$name" "$out" "${args[@]}"
}

# OOM-CHECK before the Gemma Q8 block (worst case g8+dflash-f16, -ub 512).
oom_check_g8() {
    log "OOM-CHECK: gemma q8 + dflash f16 (llama-cli -ub 512, 1 prompt n=8)"
    if "$LLAMA_CPP_BIN/llama-cli" -m "$G-Q8_0.gguf" -md "$DK/gemma-4-12B-it-DFlash-F16.gguf" \
        --spec-type draft-dflash -ngl 99 -ngld 99 -n 8 -p "Hello." -ub 512 \
        --no-conversation --single-turn --reasoning off -c 2048 -t 8 --seed 42 \
        < /dev/null > /tmp/opencode/oom-check-curves.log 2>&1; then
        log "OOM-CHECK: OK (gemma q8 usable with -ub 512)"
        OOM_OK=1
    else
        log "OOM-CHECK: FAILED → SKIP the g8 configs (see /tmp/opencode/oom-check-curves.log)"
        OOM_OK=0
    fi
}

log "CURVES CHAIN START — ${#MATRIX[@]} configs with drafter × 180 prompts (acc-sample.jsonl)"

if [ "${DRY_RUN:-0}" = "1" ]; then
    echo "DRY-RUN: curves matrix (${#MATRIX[@]}):"
    for entry in "${MATRIX[@]}"; do
        IFS='|' read -r dname dtarget ddraft dstype dpmin dout <<< "$entry"
        echo "  $dname | ${ddraft:-—} | ${dstype:-—} | ${dpmin:-—}"
    done
    echo "DRY-RUN: end (${#MATRIX[@]} configs, ~15 min each ≈ 4.5 h)"
    exit 0
fi

T0=$(date +%s)
for entry in "${MATRIX[@]}"; do
    name="${entry%%|*}"
    case "$name" in
        curves-gemma-q8-*)
            if [ "$G8_CHECKED" -eq 0 ]; then
                oom_check_g8
                G8_CHECKED=1
            fi
            if [ "$OOM_OK" != "1" ]; then
                log "SKIP: $name (OOM-CHECK g8 failed)"
                SKIP_COUNT=$((SKIP_COUNT + 1))
                continue
            fi
            ;;
    esac
    run_config "$entry"
done

DUR=$(( $(date +%s) - T0 ))
log "CURVES CHAIN DONE (duration ${DUR}s ≈ $((DUR / 3600))h $((DUR % 3600 / 60))m)"
log "CURVES SUMMARY: OK=$OK_COUNT FAILED=$FAILED_COUNT SKIP=$SKIP_COUNT"
if [ "$FAILED_COUNT" -eq 0 ] && [ "$SKIP_COUNT" -eq 0 ]; then
    log "CURVES CHAIN: exit 0"
    exit 0
fi
log "CURVES CHAIN: exit 1 (FAILED=$FAILED_COUNT SKIP=$SKIP_COUNT — check queue-curves.log)"
exit 1