#!/usr/bin/env python3
"""
Build the camstack-* model set with batch-flexible CoreML packages.

The pipeline's Python pool currently does N predicts iterated internally
on a list-of-dicts (saves session setup, but each frame still costs ~9ms
warm). A CoreML model exported with a flexible batch dim accepts a single
stacked tensor of shape `(N, 3, H, W)` — one ANE dispatch per batch,
amortising every fixed cost.

This script produces (per model in MODELS):
  exported-models/camstack-<id>.onnx     — ONNX with dynamic batch axis
  exported-models/camstack-<id>.mlpackage — CoreML with RangeDim batch (1..16)
  exported-models/camstack-<id>_openvino_model/  — OpenVINO IR with dynamic batch

The dynamic axis is SAFE to export and must not be pinned here: the same
artifact is loaded on CPU, iGPU, NPU and ANE nodes, so an export cannot know
which target it will meet. The loader does — `_static_batch_required` /
`_pin_batch_for_device` in `packages/addon-pipeline/python/inference_pool.py`
reshape the graph to a static batch of 1 for the one target that refuses dynamic
dims (the Intel NPU) and hand it over untouched everywhere else (D202). Pinning
it here would instead delete the flexible axis the CoreML batched-ANE dispatch
path depends on.

Usage:
  python scripts/build-camstack-models.py      # from the CamStack repository
  python build-camstack-models.py              # from the mirror's source/ folder

Environment: the packages in convert_lib's requirements.txt plus `ultralytics`
(which also fetches the upstream `.pt` checkpoints). INT8 needs calibration
images in CAMSTACK_CALIB_DIR (default ./calibration-images).
"""

import os
import shutil
import sys
from pathlib import Path

# convert_lib.py is the converter Model Studio runs at runtime. In the CamStack
# repository it lives in packages/addon-model-studio/python/; the published copy
# of this script (the `source/` folder of the camstack/camstack-models mirror)
# ships it as a sibling file. A sibling wins, so the SAME file runs in both places.
_HERE = Path(__file__).resolve().parent
_CONVERT_LIB_DIRS = (
    _HERE,
    _HERE.parent / 'packages' / 'addon-model-studio' / 'python',
)
_CONVERT_LIB_DIR = next(
    (d for d in _CONVERT_LIB_DIRS if (d / 'convert_lib.py').is_file()), None
)
if _CONVERT_LIB_DIR is None:
    sys.exit(
        'convert_lib.py not found next to this script or in '
        'packages/addon-model-studio/python/'
    )
sys.path.insert(0, str(_CONVERT_LIB_DIR))
from convert_lib import export_openvino, quantize_openvino_int8, validate_openvino_ir

# --- Config ------------------------------------------------------------

OUTPUT_DIR = Path("exported-models")

MODELS = [
    # (ultralytics-id, exported-name, image-size)
    ("yolov9t", "camstack-yolov9t", 640),
    ("yolov9s", "camstack-yolov9s", 640),
    ("yolo26n", "camstack-yolo26n", 640),
]

# Directory of calibration images (e.g. a COCO val subset) for INT8. INT8 is
# only produced when this dir exists + has images; otherwise it's skipped with
# a note. Override via CAMSTACK_CALIB_DIR.
CALIB_DIR = Path(os.environ.get("CAMSTACK_CALIB_DIR", "calibration-images"))

BATCH_RANGE_LOWER = 1
BATCH_RANGE_UPPER = 16
BATCH_DEFAULT = 1


# --- Helpers -----------------------------------------------------------


def get_size_mb(path: Path) -> float:
    if path.is_file():
        return path.stat().st_size / (1024 * 1024)
    return sum(f.stat().st_size for f in path.rglob("*") if f.is_file()) / (1024 * 1024)


def ensure_dir(p: Path) -> Path:
    p.mkdir(parents=True, exist_ok=True)
    return p


# --- Step 1: YOLO -> ONNX (dynamic batch) ------------------------------


def export_onnx(ultra_id: str, out_path: Path, imgsz: int) -> Path:
    """Export the Ultralytics .pt as ONNX with `dynamic=True`.

    `dynamic=True` produces an ONNX with `images` axis 0 marked dynamic
    (batch) so coremltools.convert can set RangeDim on it.
    """
    from ultralytics import YOLO

    print(f"  → loading {ultra_id}.pt (auto-download if missing)")
    model = YOLO(f"{ultra_id}.pt")
    print(f"  → exporting ONNX (imgsz={imgsz}, dynamic batch)")
    onnx_src = model.export(
        format="onnx",
        imgsz=imgsz,
        dynamic=True,
        simplify=True,
        opset=17,
    )
    if not onnx_src:
        raise RuntimeError(f"ONNX export returned no path for {ultra_id}")
    onnx_src = Path(onnx_src)
    if onnx_src != out_path:
        shutil.copy2(onnx_src, out_path)
    print(f"    ✓ {out_path.name} ({get_size_mb(out_path):.1f} MB)")
    return out_path


# --- Step 2: ONNX -> CoreML (RangeDim batch) ---------------------------


def export_coreml(ultra_id: str, out_path: Path, imgsz: int) -> Path:
    """Export the Ultralytics .pt as a CoreML mlpackage with flexible batch.

    Ultralytics' CoreML exporter handles the torch.jit.trace + coremltools
    pipeline internally. With `dynamic=True` + `batch=N`, it produces a
    model that accepts shape `(B, 3, imgsz, imgsz)` with B ∈ [1..N].
    `half=True` sets compute precision to FLOAT16 — required for ANE.
    """
    from ultralytics import YOLO

    print(f"  → loading {ultra_id}.pt for CoreML export")
    model = YOLO(f"{ultra_id}.pt")
    print(
        f"  → exporting CoreML mlpackage (imgsz={imgsz}, dynamic batch up to {BATCH_RANGE_UPPER}, half=FP16)"
    )
    src = model.export(
        format="coreml",
        imgsz=imgsz,
        dynamic=True,
        batch=BATCH_RANGE_UPPER,
        half=True,
        nms=False,
    )
    if not src:
        raise RuntimeError(f"CoreML export returned no path for {ultra_id}")
    src = Path(src)
    if out_path.exists():
        shutil.rmtree(out_path)
    if src.is_dir():
        shutil.copytree(src, out_path)
    else:
        out_path.mkdir(parents=True, exist_ok=True)
        shutil.copy2(src, out_path / src.name)
    print(f"    ✓ {out_path.name} ({get_size_mb(out_path):.1f} MB)")
    return out_path


# --- Step 3: ONNX -> OpenVINO (dynamic batch IR) -----------------------
# export_openvino, quantize_openvino_int8, validate_openvino_ir are imported
# from convert_lib.py (shared with the runtime; located above).

# --- Step 4: Verification ---------------------------------------------


def verify_coreml_batch(mlpackage_path: Path, imgsz: int) -> None:
    """Predict with N=1 then N=4 to confirm batch flexibility works."""
    import coremltools as ct
    import numpy as np

    print(f"  → loading {mlpackage_path.name} for verification")
    model = ct.models.MLModel(str(mlpackage_path))
    # Find input name
    spec = model.get_spec()
    input_name = spec.description.input[0].name

    for batch in (1, 4):
        x = np.random.rand(batch, 3, imgsz, imgsz).astype(np.float32)
        out = model.predict({input_name: x})
        out_keys = list(out.keys())
        out_shape = out[out_keys[0]].shape if hasattr(out[out_keys[0]], "shape") else "?"
        print(f"    ✓ batch={batch} → output[{out_keys[0]}].shape={out_shape}")


# --- Main --------------------------------------------------------------


def main() -> None:
    output = ensure_dir(OUTPUT_DIR)
    print(f"Output: {output.resolve()}")

    for ultra_id, name, imgsz in MODELS:
        print()
        print("=" * 70)
        print(f"Building {name}  (source: {ultra_id}.pt, imgsz={imgsz}×{imgsz})")
        print("=" * 70)

        onnx_path = output / f"{name}.onnx"
        coreml_path = output / f"{name}.mlpackage"

        try:
            export_onnx(ultra_id, onnx_path, imgsz)
        except Exception as e:
            print(f"  ✗ ONNX export failed: {e}")
            continue

        try:
            export_coreml(ultra_id, coreml_path, imgsz)
        except Exception as e:
            print(f"  ✗ CoreML conversion failed: {e}")

        # OpenVINO IR variants — all converted from the ONNX via ov.convert_model
        # (NOT the Ultralytics openvino exporter). fp32 baseline + fp16 (GPU/NPU).
        try:
            out_xml_fp32 = output / f"{name}.xml"
            export_openvino(str(onnx_path), str(out_xml_fp32), fp16=False)
            validate_openvino_ir(str(out_xml_fp32))
            out_xml_fp16 = output / f"{name}-fp16.xml"
            export_openvino(str(onnx_path), str(out_xml_fp16), fp16=True)
            validate_openvino_ir(str(out_xml_fp16))
        except Exception as e:
            print(f"  ✗ OpenVINO conversion failed: {e}")

        # INT8 — only when calibration images are present.
        if CALIB_DIR.is_dir() and any(CALIB_DIR.iterdir()):
            try:
                out_xml_int8 = output / f"{name}-int8.xml"
                quantize_openvino_int8(
                    str(onnx_path), str(out_xml_int8), str(CALIB_DIR), imgsz
                )
                validate_openvino_ir(str(out_xml_int8))
            except Exception as e:
                print(f"  ✗ OpenVINO INT8 quantization failed: {e}")
        else:
            print(f"  ⊘ INT8 skipped — no calibration images in {CALIB_DIR} (set CAMSTACK_CALIB_DIR)")

        if coreml_path.exists():
            try:
                verify_coreml_batch(coreml_path, imgsz)
            except Exception as e:
                print(f"  ✗ CoreML verification failed: {e}")

    # --- Summary ---
    print()
    print("=" * 70)
    print("Build complete. Files:")
    print("=" * 70)
    total = 0.0
    for f in sorted(output.iterdir()):
        size = get_size_mb(f)
        total += size
        print(f"  {f.name:50s} {size:8.1f} MB")
    print("  " + "─" * 60)
    print(f"  {'Total':50s} {total:8.1f} MB")
    print()
    print("Upload to HF camstack/camstack-models when ready.")


if __name__ == "__main__":
    sys.exit(main())