# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

from __future__ import annotations

import os
from functools import partial
from pathlib import Path

import numpy as np
import torch

from ultralytics.utils import ARM64, LINUX, LOGGER, WINDOWS
from ultralytics.utils.checks import check_requirements

from .base import BaseBackend


class OpenVINOBackend(BaseBackend):
    """Intel OpenVINO inference backend for Intel hardware acceleration.

    Loads and runs inference with Intel OpenVINO IR models (*_openvino_model/ directories). Supports automatic device
    selection and Intel-specific device targeting.
    """

    def load_model(self, weight: str | Path) -> None:
        """Load an Intel OpenVINO IR model from a .xml/.bin file pair or model directory.

        Args:
            weight (str | Path): Path to the .xml file or directory containing OpenVINO model files.
        """
        LOGGER.info(f"Loading {weight} for OpenVINO inference...")
        check_requirements("openvino>=2024.0.0")
        import openvino as ov

        core = ov.Core()
        if WINDOWS:  # Avoid reduced-precision CPU kernel failures without restricting native FP32 instructions
            core.set_property("CPU", {"INFERENCE_PRECISION_HINT": ov.Type.f32})
            # Windows VMs that report AMX as enabled can still fault on AMX INT8 kernels (0xc000001d), see
            # https://github.com/openvinotoolkit/openvino/issues/37076, so cap oneDNN just below AMX
            os.environ.setdefault("ONEDNN_MAX_CPU_ISA", "AVX512_CORE_FP16")
        fallback_device = "CPU" if core.available_devices == ["CPU"] else "AUTO"
        device_name = fallback_device

        if isinstance(self.device, str) and self.device.startswith("intel"):
            device_name = self.device.split(":")[1].upper()
            self.device = torch.device("cpu")
            if not any(d == device_name or d.startswith(f"{device_name}.") for d in core.available_devices):
                LOGGER.warning(f"OpenVINO device '{device_name}' not available. Using '{fallback_device}' instead.")
                device_name = fallback_device

        w = Path(weight)
        if not w.is_file():
            w = next(w.glob("*.xml"))

        ov_model = core.read_model(model=str(w), weights=w.with_suffix(".bin"))
        if ov_model.get_parameters()[0].get_layout().empty:
            ov_model.get_parameters()[0].set_layout(ov.Layout("NCHW"))

        self.apply_metadata(self.read_metadata(w))

        # OpenVINO CPU plugin crashes running INT8 models with dynamic shapes on Intel AMX CPUs (Sapphire Rapids and
        # newer), see https://github.com/openvinotoolkit/openvino/issues/37577, so run those as static models by
        # reshaping and recompiling per input shape in forward() instead
        cpuinfo = Path("/proc/cpuinfo")
        self.read_model = (
            partial(core.read_model, model=str(w), weights=w.with_suffix(".bin"))
            if device_name in {"CPU", "AUTO"}
            and ov_model.input().get_partial_shape().is_dynamic
            and any(op.get_type_name() == "FakeQuantize" for op in ov_model.get_ops())
            and LINUX
            and cpuinfo.exists()
            and "amx_int8" in cpuinfo.read_text()
            else None
        )
        if self.read_model is not None:
            self.dynamic = False  # fixed letterbox shapes so recompiles stay rare

        # Force sync inference because AsyncInferQueue can hang indefinitely on Intel and AMD CPUs, see
        # https://github.com/ultralytics/ultralytics/issues/25923.
        config = {"PERFORMANCE_HINT": "LATENCY"}
        if LINUX and ARM64 and device_name == "CPU":
            config["EXECUTION_MODE_HINT"] = ov.properties.hint.ExecutionMode.ACCURACY
            config["INFERENCE_PRECISION_HINT"] = ov.Type.f32
        if (
            self.task == "classify"
            and device_name.startswith("NPU")
            and "NPU_TURBO" in core.get_property(device_name, "SUPPORTED_PROPERTIES")
        ):
            config["NPU_TURBO"] = "YES"

        self.compile_model = partial(core.compile_model, device_name=device_name, config=config)
        self.ov_compiled_model = self.compile_model(ov_model)
        LOGGER.info(
            f"Using OpenVINO LATENCY mode for batch={self.batch} inference on "
            f"{', '.join(self.ov_compiled_model.get_property('EXECUTION_DEVICES'))}..."
        )
        self.ov = ov

    def forward(self, im: torch.Tensor) -> list[np.ndarray]:
        """Run Intel OpenVINO inference.

        Args:
            im (torch.Tensor): Input image tensor in BCHW format, normalized to [0, 1].

        Returns:
            (list[np.ndarray]): Model predictions as a list of numpy arrays, one per output layer.
        """
        im = im.cpu().numpy().astype(np.float32, copy=False)
        if self.read_model is not None and self.ov_compiled_model.input().get_partial_shape() != self.ov.PartialShape(
            im.shape
        ):
            ov_model = self.read_model()
            ov_model.reshape(list(im.shape))
            self.ov_compiled_model = self.compile_model(ov_model)

        return list(self.ov_compiled_model(im).values())
