#!/usr/bin/env bash
# Compile every sliced ONNX model and collect the resulting AXModels.
set -euo pipefail

SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
PROJECT_ROOT=$(cd "${SCRIPT_DIR}/.." && pwd)

ONNX_DIR="${PROJECT_ROOT}/compiled_subgraph_from_onnx/frontend"
OUTPUT_ROOT="${PROJECT_ROOT}/compiled_slice_quant_onnx"
LOG_ROOT="${OUTPUT_ROOT}/logs"
CONFIG_FILE="${PROJECT_ROOT}/pulsar2_configs/transformers.json"
TARGET_HW="AX650"
NPU_MODE="NPU3"
SKIP_FILES=("optimized.onnx" "optimized_quant_axmodel.onnx")

mkdir -p "${OUTPUT_ROOT}" "${LOG_ROOT}"

if [[ ! -d "${ONNX_DIR}" ]]; then
  echo "未找到切分后的 ONNX 目录: ${ONNX_DIR}" >&2
  exit 1
fi

mapfile -t ONNX_FILES < <(find "${ONNX_DIR}" -maxdepth 1 -type f -name '*.onnx' -printf '%f\n' | sort)

if [[ ${#ONNX_FILES[@]} -eq 0 ]]; then
  echo "目录 ${ONNX_DIR} 中没有 ONNX 文件" >&2
  exit 1
fi

echo "即将编译 ${#ONNX_FILES[@]} 个子模型"

for filename in "${ONNX_FILES[@]}"; do
  skip=false
  for banned in "${SKIP_FILES[@]}"; do
    if [[ "${filename}" == "${banned}" ]]; then
      skip=true
      break
    fi
  done
  if [[ "${skip}" == true ]]; then
    echo "跳过 ${filename}"
    continue
  fi

  stem="${filename%.onnx}"
  input_path="${ONNX_DIR}/${filename}"
  output_dir="${OUTPUT_ROOT}/${stem}"
  output_name="${stem}.axmodel"
  log_path="${LOG_ROOT}/${stem}.log"

  mkdir -p "${output_dir}"

  echo "[pulsar2] ${filename} -> ${output_name}"
  if ! pulsar2 build \
    --input "${input_path}" \
    --model_type QuantAxModel \
    --output_dir "${output_dir}" \
    --output_name "${output_name}" \
    --config "${CONFIG_FILE}" \
    --target_hardware "${TARGET_HW}" \
    --npu_mode "${NPU_MODE}" \
    --debug.dump_frontend_graph true \
    2>&1 | tee "${log_path}"; then
    echo "编译 ${filename} 失败，日志已保存到 ${log_path}" >&2
    continue
  fi

  cp -f "${output_dir}/${output_name}" "${OUTPUT_ROOT}/"
  echo "已复制 ${output_name} 到 ${OUTPUT_ROOT}"
done

echo "全部子模型编译完成。AXModel 文件已汇总至 ${OUTPUT_ROOT}"