#!/usr/bin/env bash
set -Eeuo pipefail

ROOT=/opt/h3
MODELS="$ROOT/models"
QUALITY="$ROOT/quality"
LOGS="$ROOT/logs/quality"
SD_CLI="$ROOT/src/stable-diffusion.cpp/build/bin/sd-cli"
Q8="$MODELS/minimax_h3_fl2va_pruned-Q8_0.gguf"
LLM="$MODELS/qwen3vl_32b_minimax_h3-Q2_K_M.gguf"
VIDEO_VAE="$MODELS/vae/minimax_h3_video_vae_fp16.safetensors"
AUDIO_VAE="$MODELS/vae/minimax_h3_audio_vae_fp32.safetensors"
IMATRIX="$QUALITY/calibration/minimax_h3_fl2va_pruned-qf.imatrix"
PUBLISHED_IMATRIX_URL='https://huggingface.co/MarxistLeninist/MiniMax-H3-FL2VA-Pruned-IQ1-GGUF/resolve/456f2691119a5cd065ccd2ca27e71dff1dd6624e/calibration/minimax_h3_fl2va_pruned.imatrix?download=true'

# First match wins. The two text-token-refiner blocks and final transformer
# blocks stay Q8. Attention and residual/down projections stay Q4_K. The
# refreshed importance matrix shows a steep rise across the later fc1 blocks,
# so blocks 30-45 are also Q4_K while the 30 lower-importance fc1 matrices in
# blocks 0-29 remain true IQ1. This keeps 4.624B parameters at IQ1 while
# protecting the matrices that dominate the measured activation importance.
QF_RULES='condition_proj\.weight$=bf16,token_refiner\.=q8_0,blocks\.(46|47|48|49)\.=q8_0,blocks\.(3[0-9]|4[0-5])\.mlp\.fc1\.weight$=q4_K,blocks\.[0-9]+\.attn\.(qkv_proj|out_proj)\.weight$=q4_K,blocks\.[0-9]+\.mlp\.fc2\.weight$=q4_K'

mkdir -p "$QUALITY/calibration" "$QUALITY/candidates" "$QUALITY/eval" "$LOGS"

require_file() {
  local path="$1"
  [[ -s "$path" ]] || { echo "required file missing or empty: $path" >&2; exit 1; }
}

require_file "$SD_CLI"
require_file "$Q8"
require_file "$LLM"
require_file "$VIDEO_VAE"
require_file "$AUDIO_VAE"

calibrate() {
  if [[ ! -s "$IMATRIX" ]]; then
    curl -fL --retry 5 --retry-all-errors -o "$IMATRIX.part" "$PUBLISHED_IMATRIX_URL"
    mv "$IMATRIX.part" "$IMATRIX"
  fi

  local prompts=(
    'A red fox trots through crunchy snow; audible footsteps, soft wind, natural synchronized stereo sound.'
    'A ginger cat surfs a blue ocean wave, spray sparkling in sunlight; rolling surf and playful meow in stereo.'
    'Close-up of a woman speaking to camera in a quiet kitchen, natural blinking and lip motion; clear conversational voice.'
    'A drummer performs on a concert stage under moving colored lights; synchronized drums and cheering crowd.'
    'Cars and bicycles cross a rainy neon city intersection at night; tires on wet road and distant traffic.'
    'A red panda steps along a mossy log in a misty forest, detailed fur and foliage; birds and gentle forest ambience.'
    'Two hands slowly pour steaming coffee into a ceramic cup beside a window; liquid sound and soft room tone.'
    'A polished silver robot flies above a desert canyon at sunset, cinematic tracking shot; wind rush and mechanical hum.'
  )
  local seeds=(11 23 37 41 53 67 79 97)
  local widths=(320 384 320 384 320 384 320 384)
  local heights=(192 256 192 256 192 256 192 256)
  local frames=(22 25 22 25 22 25 22 25)

  for i in "${!prompts[@]}"; do
    local n=$((i + 1))
    local out="$QUALITY/calibration/cal-${n}.webm"
    local log="$LOGS/calibration-${n}.log"
    "$SD_CLI" \
      --mode vid_gen \
      --diffusion-model "$Q8" \
      --llm "$LLM" \
      --vae "$VIDEO_VAE" \
      --audio-vae "$AUDIO_VAE" \
      --prompt "${prompts[$i]}" \
      --width "${widths[$i]}" --height "${heights[$i]}" \
      --video-frames "${frames[$i]}" --fps 24 \
      --steps 4 --cfg-scale 1.0 --backend te=cpu --diffusion-fa \
      --rng cpu --seed "${seeds[$i]}" \
      --imat-in "$IMATRIX" --imat-out "$IMATRIX" \
      --output "$out" --verbose >"$log" 2>&1
    ffprobe -v error -select_streams v:0 -show_entries stream=width,height,nb_frames -of default=nw=1 "$out" >>"$log" 2>&1
  done

  sha256sum "$IMATRIX" | tee "$QUALITY/calibration/IMATRIX_SHA256"
  stat -c '%n %s bytes' "$IMATRIX"
}

build_one() {
  local label="${1:?expected IQ1_M or IQ1_S}"
  local qtype
  case "$label" in
    IQ1_M) qtype=iq1_m ;;
    IQ1_S) qtype=iq1_s ;;
    *) echo "unsupported label: $label" >&2; exit 2 ;;
  esac

  require_file "$IMATRIX"
  local out="$QUALITY/candidates/minimax_h3_fl2va_pruned-UD-${label}-QF.gguf"
  local part="$out.part"
  local log="$LOGS/build-UD-${label}-QF.log"
  rm -f "$part"

  "$SD_CLI" \
    --mode convert \
    --diffusion-model "$Q8" \
    --type "$qtype" \
    --tensor-type-rules "$QF_RULES" \
    --imat-in "$IMATRIX" \
    --threads "$(nproc)" \
    --output "$part" --verbose >"$log" 2>&1

  [[ -s "$part" ]]
  mv "$part" "$out"
  sha256sum "$out" | tee "$out.sha256"
  stat -c '%n %s bytes' "$out"
}

case "${1:-}" in
  calibrate) calibrate ;;
  build)
    build_one "${2:?expected IQ1_M or IQ1_S}"
    ;;
  all)
    calibrate
    build_one IQ1_M
    build_one IQ1_S
    ;;
  *)
    echo "usage: $0 calibrate | build IQ1_M | build IQ1_S | all" >&2
    exit 2
    ;;
esac