#!/usr/bin/env bash
# Export YOLO .pt/.onnx → TensorRT .engine (FP16).
#
#   bash ~/ros2/scripts/export_yolo_trt.sh
#   bash ~/ros2/scripts/export_yolo_trt.sh /home/cc/ros2/models/yolo/detect_best.onnx 640
#
# Needs: tensorrt-cu12 + onnx + onnxslim in ros_flexiv (do NOT pip install
# the meta package `tensorrt`, it pulls CUDA 13 wheels).
# Then set yolo.backend: tensorrt in site.yaml and restart visual servo.

set -eo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
# shellcheck disable=SC1091
source "${ROOT}/scripts/lib/load_site_config.sh"

PT="${1:-${SITE_YOLO_WEIGHTS:-${ROOT}/models/yolo/detect_best.onnx}}"
IMGSZ="${2:-${SITE_YOLO_IMGSZ:-640}}"
PY="${HOME}/anaconda3/envs/ros_flexiv/bin/python"

if [[ ! -f "${PT}" ]]; then
  echo "ERROR: weights not found: ${PT}" >&2
  exit 1
fi

# Pip TensorRT shared libs (needed at export + runtime).
TRT_LIBS="$("${PY}" - <<'PY'
import pathlib, tensorrt
p = pathlib.Path(tensorrt.__file__).resolve().parent
cands = [p, p.parent / "tensorrt_libs", p.parent / "nvidia" / "tensorrt" / "lib"]
print(":".join(str(x) for x in cands if x.exists()))
PY
)"
export LD_LIBRARY_PATH="${TRT_LIBS}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
export YOLO_AUTOINSTALL=false

echo "Export TensorRT engine from ${PT}  imgsz=${IMGSZ}"
"${PY}" - <<PY
from pathlib import Path
import tensorrt as trt

src = Path("${PT}").expanduser()
print("tensorrt", trt.__version__)
print("source", src)
engine_path = src.with_suffix(".engine")

if src.suffix.lower() == ".onnx":
    logger = trt.Logger(trt.Logger.INFO)
    flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    builder = trt.Builder(logger)
    network = builder.create_network(flag)
    parser = trt.OnnxParser(network, logger)
    raw = src.read_bytes()
    if not parser.parse(raw):
        for i in range(parser.num_errors):
            print("ONNX parse error:", parser.get_error(i))
        raise SystemExit("ONNX parse failed")
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30)
    if builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
        print("FP16 enabled")
    serialized = builder.build_serialized_network(network, config)
    if serialized is None:
        raise SystemExit("TensorRT build failed")
    engine_path.write_bytes(bytes(serialized))
    print("wrote", engine_path)
else:
    from ultralytics import YOLO
    m = YOLO(str(src))
    try:
        out = m.export(
            format="engine",
            imgsz=int("${IMGSZ}"),
            device=0,
            quantize=16,
            simplify=True,
            workspace=4,
        )
    except TypeError:
        out = m.export(
            format="engine",
            imgsz=int("${IMGSZ}"),
            device=0,
            half=True,
            simplify=True,
            workspace=4,
        )
    print("wrote", out)
PY
echo "Done. Restart: bash ~/ros2/scripts/run_visual_servo_rt.sh --control"