# coding=utf-8
"""导出 YOLO / ConvNeXt 特征模型为 OpenVINO，供无 NVIDIA 的客户机加速。

用法（在项目根目录）::

    pip install openvino onnx
    python tools/export_openvino.py
    python tools/export_openvino.py --yolo-only
    python tools/export_openvino.py --feature-only
    python tools/export_openvino.py --imgsz 416

导出后：
  - best_openvino_model/     → 启动时自动被 YOLO 加载
  - feature_openvino_model/  → 启动时自动被特征提取加载

环境变量（可选）:
  DISH_YOLO_MODEL / DISH_FEATURE_OPENVINO / DISH_OPENVINO_DEVICE / DISH_YOLO_IMGSZ
"""
from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path

_ROOT = Path(__file__).resolve().parents[1]
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))


def export_yolo(imgsz: int = 640) -> Path:
    from ultralytics import YOLO

    pt = _ROOT / "best.pt"
    if not pt.exists():
        raise FileNotFoundError(f"找不到权重: {pt}")
    model = YOLO(str(pt))
    out = model.export(format="openvino", imgsz=imgsz, half=False)
    out_path = Path(out)
    # ultralytics 通常输出 best_openvino_model/
    target = _ROOT / "best_openvino_model"
    if out_path.resolve() != target.resolve():
        if target.exists():
            shutil.rmtree(target)
        if out_path.is_dir():
            shutil.move(str(out_path), str(target))
        else:
            target.mkdir(parents=True, exist_ok=True)
            shutil.copy2(out_path, target / out_path.name)
    print(f"[OK] YOLO OpenVINO → {target}")
    return target


def export_feature(imgsz: int = 224) -> Path:
    """导出 ConvNeXt-Tiny 特征骨干为 OpenVINO（输入 NCHW 已归一化）。"""
    import numpy as np
    import torch
    from openvino import convert_model, save_model

    from core.common_classifier import FeatureExtractor

    extractor = FeatureExtractor().eval()
    dummy = torch.randn(1, 3, imgsz, imgsz)

    onnx_path = _ROOT / "feature_extractor.onnx"
    torch.onnx.export(
        extractor,
        dummy,
        str(onnx_path),
        input_names=["input"],
        output_names=["features"],
        dynamic_axes={"input": {0: "batch"}, "features": {0: "batch"}},
        opset_version=17,
    )
    print(f"[OK] ONNX → {onnx_path}")

    ov_model = convert_model(str(onnx_path))
    out_dir = _ROOT / "feature_openvino_model"
    out_dir.mkdir(parents=True, exist_ok=True)
    xml_path = out_dir / "feature_extractor.xml"
    save_model(ov_model, str(xml_path))
    print(f"[OK] feature OpenVINO → {xml_path}")

    # 快速校验
    from openvino import Core
    compiled = Core().compile_model(str(xml_path), "CPU")
    test = np.random.randn(1, 3, imgsz, imgsz).astype(np.float32)
    out = compiled([test])[compiled.outputs[0]]
    print(f"[OK] smoke test output shape={np.asarray(out).shape}")
    return out_dir


def main():
    parser = argparse.ArgumentParser(description="导出 OpenVINO 模型")
    parser.add_argument("--yolo-only", action="store_true")
    parser.add_argument("--feature-only", action="store_true")
    parser.add_argument("--imgsz", type=int, default=640, help="YOLO 输入边长")
    args = parser.parse_args()

    do_yolo = not args.feature_only
    do_feat = not args.yolo_only
    if args.yolo_only:
        do_feat = False
    if args.feature_only:
        do_yolo = False

    if do_yolo:
        export_yolo(imgsz=args.imgsz)
    if do_feat:
        export_feature(imgsz=224)
    print("完成。重启识别程序即可自动加载 OpenVINO 模型。")


if __name__ == "__main__":
    main()