#!/usr/bin/env bash
# k sweep (drafter block length) — uli decision 2026-08-07:
# k ∈ {5, 7, 10} over 7 representative configs × 180 prompts (acc-sample.jsonl).
# k=3 is already covered by the curves chain (curves-*, --spec-draft-n-max 3).
#
# Goal: α(k) and αbe(k) curves (break-even per k) for G4/outline F3-T5.
# Representativeness: vanilla (draft-simple), EAGLE-3, DFlash (negative control
# Qwen), DSpark-p0 (best gating), MTP (best Gemma) and the Q8 dimension (eagle3).
#
# Usage:
#   DRY_RUN=1 bash scripts/run_ksweep_chain.sh   # prints the matrix (21) and exits
#   bash scripts/run_ksweep_chain.sh             # ~40 min/run ≈ 14 h total
#
# RUN ONLY WITH THE GPU FREE (after the curves chain and the ctx-8192 fix).
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-ksweep.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"
KS=(5 7 10)

OK_COUNT=0
FAILED_COUNT=0
SKIP_COUNT=0

# k-sweep matrix (7 configs × KS): name|target|draft|spec_type|p_min|out|extra|k
# (out = experiments/runs/ksweep-<base>-k<k>; k is expanded in run_config).
BASE_MATRIX=(
    "ksweep-qwen-q4-vanilla17b|$Q-Q4_K_M.gguf|$DK/Qwen3-1.7B-Q4_K_M.gguf|draft-simple||"
    "ksweep-qwen-q4-eagle3|$Q-Q4_K_M.gguf|$DK/Qwen3-8B-speculator.eagle3-F16.gguf|draft-eagle3||"
    "ksweep-qwen-q4-dflash|$Q-Q4_K_M.gguf|$DK/dflash_qwen3_8b_block7.gguf|draft-dflash||"
    "ksweep-qwen-q4-dspark-p0|$Q-Q4_K_M.gguf|$DK/dspark_qwen3_8b_block7.gguf|draft-dspark|0.0|"
    "ksweep-qwen-q8-eagle3|$Q-Q8_0.gguf|$DK/Qwen3-8B-speculator.eagle3-F16.gguf|draft-eagle3||"
    "ksweep-gemma-q4-mtp|$G-Q4_K_M.gguf|$DK/mtp-gemma-4-12b-it-Q8_0.gguf|draft-mtp||"
    "ksweep-gemma-q4-dflash-f16|$G-Q4_K_M.gguf|$DK/gemma-4-12B-it-DFlash-F16.gguf|draft-dflash||"
)

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

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

log "KSWEEP CHAIN START — ${#BASE_MATRIX[@]} configs × k${KS[*]} (180 prompts acc-sample)"

if [ "${DRY_RUN:-0}" = "1" ]; then
    echo "DRY-RUN: k-sweep matrix ($((${#BASE_MATRIX[@]} * ${#KS[@]})) runs):"
    for entry in "${BASE_MATRIX[@]}"; do
        IFS='|' read -r dname dtarget ddraft dstype dpmin dextra <<< "$entry"
        for k in "${KS[@]}"; do
            echo "  ${dname}-k${k} | $dtarget | ${ddraft:-—} | ${dstype:-—} | ${dpmin:-—} | --spec-draft-n-max $k"
        done
    done
    echo "DRY-RUN: end ($((${#BASE_MATRIX[@]} * ${#KS[@]})) runs, ~40 min each ≈ $(( ${#BASE_MATRIX[@]} * ${#KS[@]} * 40 / 60 ))h)"
    exit 0
fi

T0=$(date +%s)
for entry in "${BASE_MATRIX[@]}"; do
    for k in "${KS[@]}"; do
        run_config "$entry" "$k"
    done
done

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