#!/usr/bin/env bash
# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Export a YOLO26 or DETR detector to OpenVINO IR for the object-counting use case.
# Usage: ./export_and_quantize.sh [MODEL_VARIANT] [PRECISION]
# Example: ./export_and_quantize.sh yolo26n FP16
# Example: ./export_and_quantize.sh detr-resnet-50 FP16
#
# Supported model variants:
#   yolo26n, yolo26s, yolo26m, yolo26l           -- Ultralytics YOLO26 (default)
#   detr-resnet-50                               -- Facebook DETR (Hugging Face)
#
# DETR is exported end-to-end: ImageNet normalization, softmax,
# box decoding, and the COCO-91 to COCO-80 class remap are embedded in the
# graph, so the IR produces the same [1, N, 6] output layout and class ids as
# YOLO26 and runs with the same OpenVINO and DLStreamer samples.
#
# Supported precisions:
#   FP32  -- Full-precision floating-point weights
#   FP16  -- Half-precision floating-point weights (default)
#   INT8  -- Quantized 8-bit integer weights (requires NNCF)
#
# Precision / device compatibility:
#   | Precision | CPU | GPU | NPU |
#   |-----------|-----|-----|-----|
#   | FP32      | Yes | Yes | No  |
#   | FP16      | Yes | Yes | Yes |
#   | INT8      | Yes | Yes | Yes |

set -euo pipefail

MODEL_NAME="${1:-yolo26n}"
PRECISION="${2:-FP16}"
PRECISION="$(echo "${PRECISION}" | tr '[:lower:]' '[:upper:]')"

if [[ "${PRECISION}" != "FP32" && "${PRECISION}" != "FP16" && "${PRECISION}" != "INT8" ]]; then
    echo "ERROR: unsupported precision '${PRECISION}'. Choose FP32, FP16, or INT8." >&2
    exit 1
fi

case "${MODEL_NAME}" in
    yolo26n|yolo26s|yolo26m|yolo26l)
        MODEL_FAMILY="yolo"
        INPUT_SIZE=640
        ;;
    detr-resnet-50)
        MODEL_FAMILY="detr"
        INPUT_SIZE=800
        ;;
    *)
        echo "ERROR: unsupported model '${MODEL_NAME}'. Choose yolo26[n|s|m|l] or detr-resnet-50." >&2
        exit 1
        ;;
esac

# Ask for approval before downloading models and sample files
echo ""
echo "This script will download:"
echo "  - Model weights and/or sample files"
echo ""
read -p "Continue with downloads? (yes/no): " APPROVAL
if [[ "${APPROVAL}" != "yes" ]]; then
    echo "Download cancelled by user."
    exit 0
fi

# Ping the HuggingFace repo to register a tracked download of config.json.
# This is best-effort: a failure (offline, or before the repo is published)
# must not stop the export.
echo "--- Registering HuggingFace download (tracking ping) ---"
HF_REPO_ID="Intel/object-counting"
HF_CONFIG_URL="https://huggingface.co/${HF_REPO_ID}/resolve/main/config.json"
if curl -fsSL -o /dev/null "${HF_CONFIG_URL}"; then
    echo "Registered HuggingFace download for ${HF_REPO_ID}"
else
    echo "WARNING: HuggingFace tracking ping failed (offline?); continuing." >&2
fi
echo ""
echo "--- Installing dependencies ---"
if [[ "${MODEL_FAMILY}" == "detr" ]]; then
    DEPS=(openvino transformers torch pyyaml opencv-python)
else
    DEPS=(openvino ultralytics)
fi
if [[ "${PRECISION}" == "INT8" ]]; then
    DEPS+=(nncf)
fi
pip install -qU "${DEPS[@]}"

echo "--- Downloading sample test image ---"
if [[ ! -f test.jpg ]]; then
    wget -q -O test.jpg https://ultralytics.com/images/bus.jpg
    echo "Downloaded: test.jpg"
else
    echo "Already present: test.jpg"
fi
echo ""
echo "--- Downloading sample test video ---"
if [[ ! -f test_video.mp4 ]]; then
    wget -q -O test_video.mp4 \
        "https://raw.githubusercontent.com/intel-iot-devkit/sample-videos/master/person-bicycle-car-detection.mp4"
    echo "Downloaded: test_video.mp4"
else
    echo "Already present: test_video.mp4"
fi

if [[ "${PRECISION}" == "FP32" ]]; then
    HALF_FLAG="False"
    EXPORT_LABEL="FP32"
else
    HALF_FLAG="True"
    EXPORT_LABEL="FP16"
fi

if [[ "${MODEL_FAMILY}" == "yolo" ]]; then
    echo "--- Exporting ${MODEL_NAME} to OpenVINO IR (${EXPORT_LABEL}) ---"
    python3 -c "
from ultralytics import YOLO

model = YOLO('${MODEL_NAME}.pt')
model.export(format='openvino', half=${HALF_FLAG}, dynamic=False, imgsz=${INPUT_SIZE}, nms=False)
print('Export complete: ${MODEL_NAME}_openvino_model/')
"
else
    echo "--- Exporting facebook/${MODEL_NAME} to OpenVINO IR (${EXPORT_LABEL}) ---"
    python3 - "${MODEL_NAME}" "${HALF_FLAG}" "${INPUT_SIZE}" <<'PYEXPORT'
import sys
from pathlib import Path

import openvino as ov
import torch
import yaml
from transformers import DetrForObjectDetection

model_name, half_flag, size = sys.argv[1], sys.argv[2] == "True", int(sys.argv[3])
model_id = f"facebook/{model_name}"

# The "no_timm" revision uses the transformers ResNet backbone (no timm dependency).
detr = DetrForObjectDetection.from_pretrained(model_id, revision="no_timm").eval()
id2label = detr.config.id2label
# DETR predicts 91 COCO category ids (with "N/A" gaps) plus a no-object class.
# Keep the 80 real categories so class ids match YOLO26 (person = 0).
coco80 = [k for k in sorted(id2label) if id2label[k] != "N/A"]
assert len(coco80) == 80, f"unexpected label map with {len(coco80)} classes"


class DetrEndToEnd(torch.nn.Module):
    """Input: RGB [0, 1] NCHW. Output: [1, 100, 6] = [x1, y1, x2, y2, conf, class_id] in input pixels."""

    def __init__(self, model: torch.nn.Module) -> None:
        super().__init__()
        self.model = model
        self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
        self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
        self.register_buffer("coco80", torch.tensor(coco80))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.model(pixel_values=(x - self.mean) / self.std, return_dict=True)
        probs = out.logits.softmax(-1).index_select(-1, self.coco80)
        conf, cls = probs.max(-1)
        cx, cy, w, h = (out.pred_boxes * size).unbind(-1)
        return torch.stack([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2, conf, cls.float()], -1)


ov_model = ov.convert_model(
    DetrEndToEnd(detr), example_input=torch.zeros(1, 3, size, size), input=[1, 3, size, size]
)
ov_model.outputs[0].get_tensor().set_names({"output0"})

names = {i: id2label[k] for i, k in enumerate(coco80)}
# DLStreamer reads preprocessing and the post-processor type from model_info.
model_info = {
    "model_type": "YOLO",
    "labels": " ".join(n.replace(" ", "_") for n in names.values()),
    "resize_type": "standard",
    "reverse_input_channels": "YES",
    "scale_values": "255",
    "iou_threshold": "0.7",
}
for key, value in model_info.items():
    ov_model.set_rt_info(value, ["model_info", key])

out_dir = Path(f"{model_name}_openvino_model")
out_dir.mkdir(exist_ok=True)
ov.save_model(ov_model, out_dir / f"{model_name}.xml", compress_to_fp16=half_flag)
metadata = {
    "description": f"{model_id} exported end-to-end with YOLO26-compatible output",
    "license": "Apache-2.0",
    "task": "detect",
    "batch": 1,
    "imgsz": [size, size],
    "names": names,
    "channels": 3,
    "end2end": True,
}
with open(out_dir / "metadata.yaml", "w", encoding="utf-8") as f:
    yaml.safe_dump(metadata, f, sort_keys=False)
print(f"Export complete: {out_dir}/")
PYEXPORT
fi

if [[ "${PRECISION}" == "INT8" ]]; then
    echo "--- Quantizing to INT8 with NNCF ---"
    python3 -c "
import nncf
import openvino as ov
import numpy as np
import cv2

core = ov.Core()
model = core.read_model('${MODEL_NAME}_openvino_model/${MODEL_NAME}.xml')

# Use the downloaded test image for calibration instead of random noise.
img = cv2.imread('test.jpg')
img = cv2.resize(img, (${INPUT_SIZE}, ${INPUT_SIZE}))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
img = img.transpose(2, 0, 1)[np.newaxis, ...]  # NCHW

def transform_fn(data_item):
    return img

calibration_dataset = nncf.Dataset(list(range(300)), transform_fn)

quantized = nncf.quantize(
    model,
    calibration_dataset,
    preset=nncf.QuantizationPreset.MIXED,
    subset_size=300,
    model_type=nncf.ModelType.TRANSFORMER if '${MODEL_FAMILY}' == 'detr' else None,
)

ov.save_model(quantized, '${MODEL_NAME}_objcount_int8.xml')
print('Quantization complete: ${MODEL_NAME}_objcount_int8.xml')
"
fi
echo "--- Done ---"