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

from __future__ import annotations

import platform
from pathlib import Path

import numpy as np
import torch

from ultralytics.utils import LOGGER

from .base import BaseBackend, read_tflite_metadata


class TensorFlowBackend(BaseBackend):
    """Google TensorFlow inference backend supporting multiple serialization formats.

    Loads and runs inference with Google TensorFlow models in SavedModel, GraphDef (.pb), and Edge TPU formats. Handles
    quantized model dequantization and task-specific output formatting.
    """

    def __init__(self, weight: str | Path, device: torch.device, fp16: bool = False, format: str = "saved_model"):
        """Initialize the Google TensorFlow backend.

        Args:
            weight (str | Path): Path to the SavedModel directory, .pb file, or Edge TPU .tflite file.
            device (torch.device): Device to run inference on.
            fp16 (bool): Whether to use FP16 half-precision inference.
            format (str): Model format, one of "saved_model", "pb", or "edgetpu".
        """
        assert format in {"saved_model", "pb", "edgetpu"}, f"Unsupported TensorFlow format: {format}."
        self.format = format
        super().__init__(weight, device, fp16)

    def load_model(self, weight: str | Path) -> None:
        """Load a Google TensorFlow model in SavedModel, GraphDef, or Edge TPU format.

        Args:
            weight (str | Path): Path to the model file or directory.
        """
        if self.format in {"saved_model", "pb"}:
            import tensorflow as tf

        if self.format == "saved_model":
            LOGGER.info(f"Loading {weight} for TensorFlow SavedModel inference...")
            self.model = tf.saved_model.load(weight)
            # Load metadata
            metadata_file = Path(weight) / "metadata.yaml"
            if metadata_file.exists():
                from ultralytics.utils import YAML

                self.apply_metadata(YAML.load(metadata_file))
        elif self.format == "pb":
            LOGGER.info(f"Loading {weight} for TensorFlow GraphDef inference...")
            from ultralytics.utils.export.tensorflow import gd_outputs

            def wrap_frozen_graph(gd, inputs, outputs):
                """Wrap a TensorFlow frozen graph for inference by pruning to specified input/output nodes."""
                x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), [])
                ge = x.graph.as_graph_element
                return x.prune(tf.nest.map_structure(ge, inputs), tf.nest.map_structure(ge, outputs))

            gd = tf.Graph().as_graph_def()
            with open(weight, "rb") as f:
                gd.ParseFromString(f.read())
            self.frozen_func = wrap_frozen_graph(gd, inputs="x:0", outputs=gd_outputs(gd))

            # Try to find metadata
            try:
                metadata_file = next(
                    Path(weight).resolve().parent.rglob(f"{Path(weight).stem}_saved_model*/metadata.yaml")
                )
                from ultralytics.utils import YAML

                self.apply_metadata(YAML.load(metadata_file))
            except StopIteration:
                pass
        else:  # edgetpu
            try:
                from tflite_runtime.interpreter import Interpreter, load_delegate

                self.tf = None
            except ImportError:
                import tensorflow as tf

                self.tf = tf
                Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate

            device = self.device[3:] if str(self.device).startswith("tpu") else ":0"
            LOGGER.info(f"Loading {weight} on device {device[1:]} for TensorFlow Lite Edge TPU inference...")
            delegate = {"Linux": "libedgetpu.so.1", "Darwin": "libedgetpu.1.dylib", "Windows": "edgetpu.dll"}[
                platform.system()
            ]
            self.interpreter = Interpreter(
                model_path=str(weight),
                experimental_delegates=[load_delegate(delegate, options={"device": device})],
            )
            self.device = torch.device("cpu")  # Edge TPU runs on CPU from PyTorch's perspective

            self.interpreter.allocate_tensors()
            self.input_details = self.interpreter.get_input_details()
            self.output_details = self.interpreter.get_output_details()

            # Load metadata embedded in the .tflite (shared helper handles metadata.json and legacy entries)
            self.apply_metadata(read_tflite_metadata(weight))

    def forward(self, im: torch.Tensor) -> list[np.ndarray]:
        """Run Google TensorFlow inference with format-specific execution and output post-processing.

        Args:
            im (torch.Tensor): Input image tensor in BHWC format (converted from BCHW by AutoBackend).

        Returns:
            (list[np.ndarray]): Model predictions as a list of numpy arrays.
        """
        im = im.cpu().numpy()
        if self.format == "saved_model":
            y = self.model.serving_default(im)
            if not isinstance(y, list):
                y = [y]
        elif self.format == "pb":
            import tensorflow as tf

            y = self.frozen_func(x=tf.constant(im))
        else:
            h, w = im.shape[1:3]

            details = self.input_details[0]
            is_int = details["dtype"] in {np.int8, np.int16}

            if is_int:
                scale, zero_point = details["quantization"]
                im = (im / scale + zero_point).astype(details["dtype"])

            self.interpreter.set_tensor(details["index"], im)
            self.interpreter.invoke()

            y = []
            for output in self.output_details:
                x = self.interpreter.get_tensor(output["index"])
                if self.task == "semantic" and x.ndim == 3:
                    # Baked argmax class map [B, H, W] of integer class IDs, not boxes or quantized logits:
                    # skip dequantization and xywh denormalization, which would corrupt and overflow the indices.
                    y.append(x)
                    continue
                if is_int:
                    scale, zero_point = output["quantization"]
                    x = (x.astype(np.float32) - zero_point) * scale
                if x.ndim == 3:
                    # Denormalize xywh by image size
                    if x.shape[-1] == 6 or self.end2end:
                        x[:, :, [0, 2]] *= w
                        x[:, :, [1, 3]] *= h
                        if self.task == "pose":
                            x[:, :, 6::3] *= w
                            x[:, :, 7::3] *= h
                    else:
                        x[:, [0, 2]] *= w
                        x[:, [1, 3]] *= h
                        if self.task == "pose":
                            x[:, 5::3] *= w
                            x[:, 6::3] *= h
                y.append(x)

        if self.task == "segment":  # segment with (det, proto) output order reversed
            if len(y[1].shape) != 4:
                y = list(reversed(y))  # should be y = (1, 116, 8400), (1, 160, 160, 32)
            if y[1].shape[-1] == 6:  # end-to-end model
                y = [y[1]]
            else:
                y[1] = np.transpose(y[1], (0, 3, 1, 2))  # should be y = (1, 116, 8400), (1, 32, 160, 160)
        elif self.task in {"semantic", "depth"} and len(y) == 1 and y[0].ndim == 4:
            y[0] = np.transpose(y[0], (0, 3, 1, 2))  # NHWC → NCHW for semantic segmentation and depth anything logits
        return [x if isinstance(x, np.ndarray) else x.numpy() for x in y]
