#!/usr/bin/env bash
# SmartTest.sh - generate the same cat+snake image with every built variant and write a speed report.
#
#   ./SmartTest.sh                    # all rows of the matrix, 1 repeat
#   REPEAT=3 ./SmartTest.sh           # median of 3 runs per row
#   ./SmartTest.sh --only Q8          # only rows whose name contains "Q8"
#
# Output (in results/):
#   results.csv      one line per run
#   results.md       comparison table (medians), ready for the model card
#   hardware.txt     machine configuration used for the run
#   images/*.png     generated 512x512 images
#   logs/*.log       raw sd-cli logs
set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
WORK="${WORK:-$PWD/work}"
BIN="$WORK/bin"
MODELS="$WORK/models"
OUT="${OUT:-$ROOT/results}"
PY="$WORK/venv/bin/python"; [[ -x "$PY" ]] || PY="$(command -v python3)"

THREADS="${THREADS:-$(nproc)}"
CPUSET="${CPUSET:-0-$((THREADS - 1))}"
REPEAT="${REPEAT:-1}"
ONLY="${ONLY:-}"   # env ONLY=<substring>, or argument: --only <substring>
[[ "${1:-}" == "--only" ]] && ONLY="${2:?--only needs a substring}"

SEED=42
STEPS=8
WIDTH=512
HEIGHT=512
PROMPT="cute cartoon cat and cute cartoon snake, best friends, hugging together, kawaii, simple clean background, children's book illustration, soft colors, adorable"
NEGATIVE="blurry, low quality, realistic, photo"

# name | variant | model key | extra env (- = none)
MATRIX=(
    "stock-Q4_K|stock|Q4_K|-"
    "avx2-Q4_K|avx2|Q4_K|-"
    "avx512-Q4_K|avx512|Q4_K|-"
    "amx_lto-Q4_K|amx_lto|Q4_K|-"
    "amx_lto-Q4_K+AMXw|amx_lto|Q4_K|SD_CPU_EXTRA_BUFT=1"
    "amx_lto-Q4_K-ffn13|amx_lto|Q4_K-ffn13|-"
    "amx_lto-Q4_K-ffn13+AMXw|amx_lto|Q4_K-ffn13|SD_CPU_EXTRA_BUFT=1"
    "amx_lto-Q4_0|amx_lto|Q4_0|-"
    "amx_lto-Q8_0|amx_lto|Q8_0|-"
    "amx_lto-Q8_0+AMXw|amx_lto|Q8_0|SD_CPU_EXTRA_BUFT=1"
    # stage 1 (VAE decode): llamafile sgemm build, with and without conv-direct for the VAE
    "llamafile-Q4_K-ffn13+AMXw|llamafile|Q4_K-ffn13|SD_CPU_EXTRA_BUFT=1|"
    "llamafile-Q4_K-ffn13+AMXw+convdirect|llamafile|Q4_K-ffn13|SD_CPU_EXTRA_BUFT=1|--conv-direct auto,vae=on"
    # stage P2: refiner weights requantized Q8_0 -> Q4_K (speed-first, opt-in)
    "llamafile-Q4_K-ffn13-rq+AMXw+convdirect|llamafile|Q4_K-ffn13-rq|SD_CPU_EXTRA_BUFT=1|--conv-direct auto,vae=on"
)

model_file() {
    case "$1" in
        Q4_K)       echo "$MODELS/z_image_turbo-Q4_K.gguf" ;;
        Q4_K-ffn13) echo "$MODELS/z_image_turbo-Q4_K-ffn13.gguf" ;;
        Q4_K-ffn13-rq) echo "$MODELS/z_image_turbo-Q4_K-ffn13-rq.gguf" ;;
        Q4_0)       echo "$MODELS/z_image_turbo-Q4_0.gguf" ;;
        Q8_0)       echo "$MODELS/z_image_turbo-Q8_0.gguf" ;;
        *) return 1 ;;
    esac
}

log() { printf '\033[1;32m[smarttest]\033[0m %s\n' "$*"; }

mkdir -p "$OUT/images" "$OUT/logs"
CSV="$OUT/results.csv"
# APPEND=1 keeps existing rows in results.csv (additive testing: only new rows are run).
if [[ "${APPEND:-0}" != "1" || ! -s "$CSV" ]]; then
    echo "name,variant,model,env,rep,status,sampling_s,generate_s,wall_s,peak_rss_gib,diff_mean_vs_ref,diff_max_vs_ref,image" > "$CSV"
fi

# ------------------------------------------------------------ hardware report
{
    echo "date:        $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
    echo "os:          $(. /etc/os-release && echo "$PRETTY_NAME")"
    echo "kernel:      $(uname -r)"
    echo "cpu:         $(lscpu | awk -F: '/^Model name/ {sub(/^ +/,"",$2); print $2}')"
    echo "sockets:     $(lscpu | awk -F: '/^Socket/ {sub(/^ +/,"",$2); print $2}')"
    echo "cores/threads per socket: $(lscpu | awk -F: '/^Core\(s\) per socket/ {sub(/^ +/,"",$2); print $2}') / $(lscpu | awk -F: '/^Thread\(s\) per core/ {sub(/^ +/,"",$2); print $2}') SMT"
    echo "numa nodes:  $(lscpu | awk -F: '/^NUMA node\(s\)/ {sub(/^ +/,"",$2); print $2}')"
    echo "cpu flags:   $(grep -o -E 'avx512f|avx512_vnni|avx512_bf16|amx_int8|amx_bf16|avx2' /proc/cpuinfo | sort -u | tr '\n' ' ')"
    echo "cgroup cpu quota: $(cat /sys/fs/cgroup/cpu.max 2>/dev/null || echo n/a)"
    echo "threads used: $THREADS   cpuset: $CPUSET"
    echo "memory:      $(free -g | awk '/^Mem:/ {print $2" GiB total"}')"
    echo "gpu:         $(nvidia-smi -L 2>&1 | head -1)  (no GPU backend used; CPU-only inference)"
    echo "gcc:         $(gcc -dumpfullversion 2>/dev/null || echo n/a)"
    echo "cmake:       $(cmake --version 2>/dev/null | head -1 || echo n/a)"
    echo "sd.cpp:      $(git -C "$WORK/stable-diffusion.cpp" rev-parse HEAD 2>/dev/null || echo n/a)"
    echo "openvino env: $("$WORK/venv-ov/bin/python" -c 'import openvino, optimum.intel, diffusers; print(f"openvino {openvino.__version__}, optimum-intel {optimum.intel.__version__}, diffusers {diffusers.__version__}")' 2>/dev/null || echo n/a)"
    echo "openvino model: HelloSun/Z-Image-Turbo-OpenVINO-INT4 (transformer/text_encoder INT4, VAE INT8)"
    echo "prompt:      $PROMPT"
    echo "negative:    $NEGATIVE"
    echo "size/steps/seed/cfg: ${WIDTH}x${HEIGHT} / $STEPS / $SEED / 1.0 (euler, simple scheduler)"
} > "$OUT/hardware.txt"

# ------------------------------------------------------------ run cases
REF_IMAGE=""
[[ -f "$OUT/images/stock-Q4_K-r1.png" ]] && REF_IMAGE="$OUT/images/stock-Q4_K-r1.png"
OV_DIR="$MODELS/Z-Image-Turbo-OpenVINO-INT4"
OV_PY="$WORK/venv-ov/bin/python"

# run_case <name> <variant> <model-label> <env-spec> <rep-command...>
# The rep-command is executed under taskset; "__IMG__" is replaced by the output image path.
# It must print sd-cli style timing lines (sampling / generate_image) on stderr.
run_case() {
    local name="$1" variant="$2" mlabel="$3" envspec="$4"; shift 4
    local extra=()
    if [[ -n "${RUN_EXTRA:-}" ]]; then read -ra extra <<< "$RUN_EXTRA"; fi
    local envprefix=()
    [[ "$envspec" != "-" ]] && envprefix=(env "$envspec")
    for rep in $(seq 1 "$REPEAT"); do
        local tag="${name}-r${rep}"
        local log_file="$OUT/logs/$tag.log"
        local img="$OUT/images/$tag.png"
        log "run $tag"
        local cmd=() a
        for a in "$@" "${extra[@]}"; do
            if [[ "$a" == "__IMG__" ]]; then cmd+=("$img"); else cmd+=("$a"); fi
        done
        # stdout and stderr of the generator both go to the log; measure_rss appends its summary line too.
        set +e
        "${envprefix[@]}" "$PY" "$ROOT/tools/measure_rss.py" taskset -c "$CPUSET" "${cmd[@]}" \
            > "$log_file" 2>&1
        local rc=$?
        set -e
        local rss_line sampling generate wall
        rss_line="$(grep -oE 'PEAK_RSS=[0-9.]+' "$log_file" | cut -d= -f2 | tail -1 || true)"
        sampling="$(grep -oE 'sampling completed, taking [0-9.]+' "$log_file" | grep -oE '[0-9.]+' | tail -1 || true)"
        generate="$(grep -oE 'generate_image completed in [0-9.]+' "$log_file" | grep -oE '[0-9.]+' | tail -1 || true)"
        wall="$(grep -oE 'WALL=[0-9.]+' "$log_file" | cut -d= -f2 | tail -1 || true)"

        local status="ok" diff_mean="" diff_max=""
        if [[ $rc -ne 0 || ! -s "$img" ]]; then
            status="failed(rc=$rc)"
        else
            [[ -z "$REF_IMAGE" ]] && REF_IMAGE="$img"
            read -r diff_mean diff_max < <("$PY" - "$REF_IMAGE" "$img" <<'PY'
import sys
import numpy as np
from PIL import Image
a = np.asarray(Image.open(sys.argv[1]).convert("RGB")).astype(int)
b = np.asarray(Image.open(sys.argv[2]).convert("RGB")).astype(int)
d = np.abs(a - b)
print(f"{d.mean():.4f} {d.max()}")
PY
)
        fi
        echo "$name,$variant,$mlabel,$envspec,$rep,$status,$sampling,$generate,$wall,$rss_line,$diff_mean,$diff_max,images/$tag.png" >> "$CSV"
        log "  $status sampling=${sampling:-?}s generate=${generate:-?}s rss=${rss_line:-?}GiB diff_mean=${diff_mean:-n/a}"
    done
}

# ---- sd.cpp matrix
for row in "${MATRIX[@]}"; do
    IFS='|' read -r name variant mkey envspec extra_args <<< "$row"
    if [[ -n "$ONLY" && "$name" != *"$ONLY"* ]]; then continue; fi

    sdbin="$BIN/$variant/sd-cli"
    model="$(model_file "$mkey")"
    if [[ ! -x "$sdbin" || ! -f "$model" ]]; then
        log "skip $name (missing binary or model)"
        echo "$name,$variant,$mkey,$envspec,0,missing,,,,,,,," >> "$CSV"
        continue
    fi
    RUN_EXTRA="${extra_args:-}"
    run_case "$name" "$variant" "$mkey" "$envspec" \
        "$sdbin" \
        --diffusion-model "$model" \
        --vae "$MODELS/ae.safetensors" \
        --llm "$MODELS/Qwen3-4B-Instruct-2507-Q4_K_M.gguf" \
        -p "$PROMPT" -n "$NEGATIVE" \
        -W "$WIDTH" -H "$HEIGHT" --steps "$STEPS" --cfg-scale 1.0 \
        --seed "$SEED" --sampling-method euler --scheduler simple \
        -t "$THREADS" -o "__IMG__"
done

# ---- OpenVINO INT4 (HelloSun/Z-Image-Turbo-OpenVINO-INT4), 8 steps to match sd.cpp; same prompt/seed/size
for ov_row in "ov-int4-8step|8"; do
    IFS='|' read -r name steps <<< "$ov_row"
    if [[ -n "$ONLY" && "$name" != *"$ONLY"* ]]; then continue; fi
    if [[ ! -x "$OV_PY" || ! -d "$OV_DIR" ]]; then
        log "skip $name (missing OpenVINO venv or model; see build.sh)"
        echo "$name,openvino,OV-INT4,-,0,missing,,,,,,,," >> "$CSV"
        continue
    fi
    run_case "$name" "openvino" "OV-INT4" "-" \
        "$OV_PY" "$ROOT/tools/ov_bench.py" \
        --model-dir "$OV_DIR" -p "$PROMPT" \
        -W "$WIDTH" -H "$HEIGHT" --steps "$steps" --seed "$SEED" -o "__IMG__"
done

# ------------------------------------------------------------ report (medians)
"$PY" - "$CSV" "$OUT/results.md" <<'PY'
import csv, statistics, sys
from collections import OrderedDict

rows = list(csv.DictReader(open(sys.argv[1])))
groups = OrderedDict()
for r in rows:
    if r["status"] != "ok":
        continue
    groups.setdefault(r["name"], []).append(r)

def med(values):
    vals = [float(v) for v in values if v not in ("", None)]
    return statistics.median(vals) if vals else None

ref = None
if "stock-Q4_K" in groups:
    ref = med([r["generate_s"] for r in groups["stock-Q4_K"]])

lines = [
    "| # | 版本 (variant) | 模型 | 額外環境變數 | 採樣 (s) | 總生成 (s) | 相對 stock | 峰值 RSS (GiB) | 與 stock 像素差 (mean/max) |",
    "|---|---|---|---|---|---|---|---|---|",
]
for i, (name, rs) in enumerate(groups.items(), 1):
    samp = med([r["sampling_s"] for r in rs])
    gen = med([r["generate_s"] for r in rs])
    rss = med([r["peak_rss_gib"] for r in rs])
    dm = med([r["diff_mean_vs_ref"] for r in rs])
    dmx = med([r["diff_max_vs_ref"] for r in rs])
    f1 = lambda v: "n/a" if v is None else f"{v:.1f}"
    f2 = lambda v: "n/a" if v is None else f"{v:.2f}"
    speed = f"{ref / gen:.2f}x" if ref and gen else "n/a"
    diff = "ref" if name == "stock-Q4_K" else (f"{dm:.2f} / {dmx:.0f}" if dm is not None else "n/a")
    env = rs[0]["env"] if rs[0]["env"] != "-" else "-"
    if "convdirect" in name:
        env = (env + " + conv-direct") if env != "-" else "conv-direct"
    lines.append(f"| {i} | {rs[0]['variant']} | {rs[0]['model']} | {env} | {f1(samp)} | {f1(gen)} | {speed} | {f2(rss)} | {diff} |")
lines.append("")
lines.append(f"n = {max((len(v) for v in groups.values()), default=0)} run(s) per row (median). Speedup = stock-Q4_K generate time / row generate time.")
open(sys.argv[2], "w").write("\n".join(lines) + "\n")
print("\n".join(lines))
PY

log "results: $CSV"
log "table:   $OUT/results.md"