"""
dheeren.inference_backend
─────────────────────────────────────────────────────────────────────────────
Auto-selects the fastest available inference backend:

  Priority:  OpenVINO  >  ONNX Runtime  >  PyTorch (fallback)

User code never needs to touch this — Diffusion.from_pretrained() picks
the best backend automatically and prints which one it's using.
"""

from pathlib import Path
import numpy as np
import torch


class _BaseBackend:
    """Common interface all backends implement."""
    def encode_text(self, token_ids: np.ndarray) -> np.ndarray: ...
    def run_unet(self, latent: np.ndarray,
                 timestep: np.ndarray,
                 context: np.ndarray) -> np.ndarray: ...
    def decode_vae(self, latent: np.ndarray) -> np.ndarray: ...
    @property
    def name(self) -> str: ...


# ─────────────────────────────────────────────────────────────────────────────
# OpenVINO backend  (fastest on Intel CPUs — uses native ISA optimizations)
# ─────────────────────────────────────────────────────────────────────────────

class OpenVINOBackend(_BaseBackend):
    def __init__(self, weights_dir: Path):
        import openvino as ov
        core = ov.Core()

        # Enable CPU threading optimizations
        core.set_property("CPU", {
            "PERFORMANCE_HINT":         "LATENCY",
            "INFERENCE_NUM_THREADS":    "0",   # 0 = use all cores
            "ENABLE_CPU_PINNING":       "YES",
        })

        def _load(name):
            xml = weights_dir / f"{name}.xml"
            model = core.read_model(str(xml))
            return core.compile_model(model, "CPU")

        self._te  = _load("text_encoder")
        self._un  = _load("unet")
        self._vae = _load("vae_decoder")

    @property
    def name(self): return "OpenVINO (Intel optimized)"

    def encode_text(self, token_ids):
        return self._te([token_ids])[0]

    def run_unet(self, latent, timestep, context):
        return self._un([latent, timestep, context])[0]

    def decode_vae(self, latent):
        return self._vae([latent])[0]


# ─────────────────────────────────────────────────────────────────────────────
# ONNX Runtime backend  (fast, works on any CPU)
# ─────────────────────────────────────────────────────────────────────────────

class ONNXBackend(_BaseBackend):
    def __init__(self, weights_dir: Path):
        import onnxruntime as ort

        opts = ort.SessionOptions()
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        opts.inter_op_num_threads = 0   # use all cores
        opts.intra_op_num_threads = 0

        def _sess(name):
            return ort.InferenceSession(
                str(weights_dir / f"{name}.onnx"),
                sess_options = opts,
                providers    = ["CPUExecutionProvider"],
            )

        self._te  = _sess("text_encoder")
        self._un  = _sess("unet")
        self._vae = _sess("vae_decoder")

    @property
    def name(self): return "ONNX Runtime (optimized CPU)"

    def encode_text(self, token_ids):
        return self._te.run(None, {"token_ids": token_ids})[0]

    def run_unet(self, latent, timestep, context):
        return self._un.run(None, {
            "latent":   latent,
            "timestep": timestep,
            "context":  context,
        })[0]

    def decode_vae(self, latent):
        return self._vae.run(None, {"latent": latent})[0]


# ─────────────────────────────────────────────────────────────────────────────
# Pure PyTorch backend  (fallback — always works, no extra installs)
# ─────────────────────────────────────────────────────────────────────────────

class PyTorchBackend(_BaseBackend):
    def __init__(self, te, unet, vae):
        self._te   = te
        self._unet = unet
        self._vae  = vae

    @property
    def name(self): return "PyTorch CPU (fallback)"

    def encode_text(self, token_ids):
        ids = torch.from_numpy(token_ids)
        with torch.no_grad():
            return self._te(ids).numpy()

    def run_unet(self, latent, timestep, context):
        with torch.no_grad():
            return self._unet(
                torch.from_numpy(latent),
                torch.from_numpy(timestep),
                torch.from_numpy(context),
            ).numpy()

    def decode_vae(self, latent):
        with torch.no_grad():
            return self._vae.decode(torch.from_numpy(latent)).numpy()


# ─────────────────────────────────────────────────────────────────────────────
# Auto-selector
# ─────────────────────────────────────────────────────────────────────────────

def create_backend(weights_dir: Path,
                   te=None, unet=None, vae=None) -> _BaseBackend:
    """
    Try backends in order of performance.
    Falls back gracefully if a backend's package isn't installed.
    """
    # OpenVINO — best on Intel
    try:
        if (weights_dir / "unet.xml").exists():
            backend = OpenVINOBackend(weights_dir)
            print(f"  Backend: {backend.name}")
            return backend
    except Exception as e:
        print(f"  OpenVINO unavailable ({e}) — trying ONNX Runtime...")

    # ONNX Runtime
    try:
        if (weights_dir / "unet.onnx").exists():
            backend = ONNXBackend(weights_dir)
            print(f"  Backend: {backend.name}")
            return backend
    except Exception as e:
        print(f"  ONNX Runtime unavailable ({e}) — falling back to PyTorch...")

    # Pure PyTorch fallback
    backend = PyTorchBackend(te, unet, vae)
    print(f"  Backend: {backend.name}")
    return backend