#!/bin/bash
#SBATCH -J bd-base-sweep
#SBATCH -o /home/l/liyj/shiying/hku_diffusion_dllm/runs/basecode/%x_%j.out
#SBATCH -e /home/l/liyj/shiying/hku_diffusion_dllm/runs/basecode/%x_%j.err
#SBATCH -N 1
#SBATCH --partition=normal
#SBATCH --cpus-per-task=2
#SBATCH --mem=4G
#SBATCH -t 00:20:00
#SBATCH --open-mode=append

set -euo pipefail

PROJECT=/home/l/liyj/shiying/hku_diffusion_dllm
cd "$PROJECT"

RUN_DIR="$PROJECT/runs/basecode"
mkdir -p "$RUN_DIR"

OUT_JSONL="$RUN_DIR/sweep_${SLURM_JOB_ID}.jsonl"
LATEST="$RUN_DIR/latest_sweep.jsonl"

echo "[info] started_at=$(date -Iseconds)"
echo "[info] node=$(hostname)"
echo "[info] output=$OUT_JSONL"

run_one() {
  local method="$1"
  local steps="$2"
  local block_size="$3"
  local num_blocks="$4"
  local remask_ratio="$5"
  local use_cache="$6"
  local tmp
  tmp="$(mktemp "$RUN_DIR/tmp_${method}_${steps}_${block_size}_${num_blocks}_XXXX.json")"
  local args=(python3 -m basecode.blockdiff.cli --method "$method" --steps "$steps" --block-size "$block_size" --num-blocks "$num_blocks" --remask-ratio "$remask_ratio" --output "$tmp")
  if [ "$use_cache" = "true" ]; then
    args+=(--use-cache)
  fi
  "${args[@]}" >/dev/null
  python3 - "$tmp" "$OUT_JSONL" "$method" "$steps" "$block_size" "$num_blocks" "$remask_ratio" "$use_cache" <<'PY'
import json
import sys
from pathlib import Path

tmp, out, method, steps, block_size, num_blocks, remask_ratio, use_cache = sys.argv[1:]
payload = json.loads(Path(tmp).read_text())
payload["sweep"] = {
    "method": method,
    "steps": int(steps),
    "block_size": int(block_size),
    "num_blocks": int(num_blocks),
    "remask_ratio": float(remask_ratio),
    "use_cache": use_cache == "true",
}
with open(out, "a", encoding="utf-8") as f:
    f.write(json.dumps(payload, ensure_ascii=False) + "\n")
Path(tmp).unlink(missing_ok=True)
PY
}

for block_size in 8 16 32; do
  run_one confidence 6 "$block_size" 1 0.5 false
  run_one dmax 4 "$block_size" 1 0.25 false
  run_one speculative 6 "$block_size" 1 0.5 true
done

for num_blocks in 2 4 8; do
  run_one multiblock 8 16 "$num_blocks" 0.5 true
done

ln -sfn "$(basename "$OUT_JSONL")" "$LATEST"
echo "[info] latest=$LATEST"
echo "[info] finished_at=$(date -Iseconds)"