#!/usr/bin/env bash
# DiffusionGemma smoke test — steel141 3090 Ti
#
# DiffusionGemma (google/diffusiongemma-26B-A4B-it, released 2026-06-10) is a
# text-DIFFUSION MoE model. It does NOT run in Ollama or stock llama.cpp:
# the standard llama-cli/llama-server reject arch 'diffusion-gemma'. It needs
# the dedicated `llama-diffusion-cli` binary from ggml-org/llama.cpp PR #24423
# (danielhanchen / Unsloth), which denoises 256-token canvas blocks in parallel.
#
# This harness runs a few non-interactive prompts and captures wall-clock + the
# diffusion step/block telemetry the CLI prints.
set -euo pipefail

CLI="${CLI:-/mnt/ai_data/diffusiongemma/llama.cpp/build/bin/llama-diffusion-cli}"
MODEL="${MODEL:-/mnt/ai_data/diffusiongemma/gguf/diffusiongemma-26B-A4B-it-Q4_K_M.gguf}"
NGL="${NGL:-99}"          # offload all layers to GPU
NPRED="${NPRED:-256}"     # one full canvas
OUTDIR="${OUTDIR:-$(dirname "$0")/results}"
mkdir -p "$OUTDIR"

[ -x "$CLI" ]   || { echo "missing CLI: $CLI" >&2; exit 1; }
[ -f "$MODEL" ] || { echo "missing model: $MODEL" >&2; exit 1; }

run() {
  local name="$1"; local prompt="$2"; shift 2
  local log="$OUTDIR/${name}.log"
  echo "=== $name ===" | tee "$log"
  echo "prompt: $prompt" | tee -a "$log"
  local t0 t1
  t0=$(date +%s.%N)
  "$CLI" -m "$MODEL" -ngl "$NGL" -n "$NPRED" --diffusion-eb auto -p "$prompt" "$@" 2>&1 | tee -a "$log"
  t1=$(date +%s.%N)
  echo "WALL_SECONDS=$(echo "$t1 - $t0" | bc)" | tee -a "$log"
  echo | tee -a "$log"
}

# 1) Plain reasoning — sanity + coherence
run reasoning "Explain in three sentences why diffusion language models can be faster than autoregressive ones."

# 2) Code — structured output the diffusion canvas has to fill coherently
run code "Write a Python function is_prime(n) with a docstring. Output only the code."

# 3) Instruction following with a hard format constraint
run format "List exactly five primary colors of light, one per line, no extra text."

echo "All runs complete. Logs in $OUTDIR"