"""CoreML model exporter.

Converts a PyTorch model to Apple CoreML ``.mlpackage`` format via
``coremltools``.  Supports FP32, FP16 (linear) and INT8 (kmeans_lut)
quantisation on macOS.  Optionally wraps the model with an iOS-compatible
NMS post-processing pipeline.
"""

from __future__ import annotations

import logging
import platform
import warnings
from pathlib import Path
from typing import Any

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import torch

from matrice_export.formats.base import BaseExporter

logger = logging.getLogger(__name__)

MACOS = platform.system() == "Darwin"


# --------------------------------------------------------------------------- #
# iOS NMS wrapper (lazy torch import)
# --------------------------------------------------------------------------- #

def _make_ios_model(model: Any, im: Any) -> Any:
    """Create an iOS NMS wrapper model.

    The ``torch.nn.Module`` subclass is defined inside this function so that
    ``import torch`` is deferred until actually needed (only when ``nms=True``).
    """
    import torch

    class _iOSModel(torch.nn.Module):
        """Normalises detection outputs for iOS CoreML NMS pipelines."""

        def __init__(self, model: torch.nn.Module, im: torch.Tensor) -> None:
            super().__init__()
            b, c, h, w = im.shape
            self.model = model
            self.nc: int = model.nc
            if w == h:
                self.normalize: float | torch.Tensor = 1.0 / w
            else:
                self.normalize = torch.tensor([1.0 / w, 1.0 / h, 1.0 / w, 1.0 / h])

        def forward(self, x: torch.Tensor):
            xywh, conf, cls = self.model(x)[0].squeeze().split((4, 1, self.nc), 1)
            return cls * conf, xywh * self.normalize

    return _iOSModel(model, im)


# --------------------------------------------------------------------------- #
# NMS pipeline builder
# --------------------------------------------------------------------------- #

def _build_nms_pipeline(
    ct_model: Any,
    im: torch.Tensor,
    names: dict[int, str] | None,
    y: Any,
    output_path: Path,
) -> str:
    """Wrap a CoreML model in an NMS pipeline and save as ``.mlmodel``.

    This mirrors the original ``pipeline_coreml()`` function.  The pipeline
    attaches an NMS post-processing layer to the detection model, producing
    ``confidence`` and ``coordinates`` outputs directly.

    Args:
        ct_model: Converted CoreML model (``coremltools.models.MLModel``).
        im: Sample input tensor (BCHW).
        names: Mapping of class-index to class-name.
        y: Raw PyTorch model output (used to infer output shapes on
            non-macOS platforms).
        output_path: Path stem used to derive the ``.mlmodel`` save path.

    Returns:
        Absolute path to the saved pipeline model.
    """
    import coremltools as ct

    batch_size, ch, h, w = list(im.shape)
    spec = ct_model.get_spec()
    out0, out1 = iter(spec.description.output)

    # Determine output shapes -------------------------------------------
    try:
        from PIL import Image

        img = Image.new("RGB", (w, h))
        out = ct_model.predict({"image": img})
        out0_shape = out[out0.name].shape
        out1_shape = out[out1.name].shape
    except Exception:
        # On Linux / Windows the CoreML runtime is unavailable -- derive
        # shapes from the PyTorch forward pass instead.
        s = tuple(y[0].shape)
        out0_shape = (s[1], s[2] - 5)
        out1_shape = (s[1], 4)

    nx = spec.description.input[0].type.imageType.width
    ny = spec.description.input[0].type.imageType.height
    na, nc = out0_shape

    if names is not None and len(names) != nc:
        logger.warning(
            "CoreML pipeline: %d names provided but model has nc=%d classes.",
            len(names),
            nc,
        )

    # Assign missing output shapes
    out0.type.multiArrayType.shape[:] = out0_shape
    out1.type.multiArrayType.shape[:] = out1_shape

    # Rebuild model from updated spec
    ct_model = ct.models.MLModel(spec)

    # -- NMS protobuf ---------------------------------------------------
    nms_spec = ct.proto.Model_pb2.Model()
    nms_spec.specificationVersion = 5
    for i in range(2):
        decoder_output = ct_model._spec.description.output[i].SerializeToString()
        nms_spec.description.input.add()
        nms_spec.description.input[i].ParseFromString(decoder_output)
        nms_spec.description.output.add()
        nms_spec.description.output[i].ParseFromString(decoder_output)

    nms_spec.description.output[0].name = "confidence"
    nms_spec.description.output[1].name = "coordinates"

    output_sizes = [nc, 4]
    for i in range(2):
        ma_type = nms_spec.description.output[i].type.multiArrayType
        ma_type.shapeRange.sizeRanges.add()
        ma_type.shapeRange.sizeRanges[0].lowerBound = 0
        ma_type.shapeRange.sizeRanges[0].upperBound = -1
        ma_type.shapeRange.sizeRanges.add()
        ma_type.shapeRange.sizeRanges[1].lowerBound = output_sizes[i]
        ma_type.shapeRange.sizeRanges[1].upperBound = output_sizes[i]
        del ma_type.shape[:]

    nms = nms_spec.nonMaximumSuppression
    nms.confidenceInputFeatureName = out0.name
    nms.coordinatesInputFeatureName = out1.name
    nms.confidenceOutputFeatureName = "confidence"
    nms.coordinatesOutputFeatureName = "coordinates"
    nms.iouThresholdInputFeatureName = "iouThreshold"
    nms.confidenceThresholdInputFeatureName = "confidenceThreshold"
    nms.iouThreshold = 0.45
    nms.confidenceThreshold = 0.25
    nms.pickTop.perClass = True
    if names is not None:
        nms.stringClassLabels.vector.extend(
            names.values() if isinstance(names, dict) else names
        )
    nms_model = ct.models.MLModel(nms_spec)

    # -- Pipeline -------------------------------------------------------
    pipeline = ct.models.pipeline.Pipeline(
        input_features=[
            ("image", ct.models.datatypes.Array(3, ny, nx)),
            ("iouThreshold", ct.models.datatypes.Double()),
            ("confidenceThreshold", ct.models.datatypes.Double()),
        ],
        output_features=["confidence", "coordinates"],
    )
    pipeline.add_model(ct_model)
    pipeline.add_model(nms_model)

    # Correct datatypes
    pipeline.spec.description.input[0].ParseFromString(
        ct_model._spec.description.input[0].SerializeToString()
    )
    pipeline.spec.description.output[0].ParseFromString(
        nms_model._spec.description.output[0].SerializeToString()
    )
    pipeline.spec.description.output[1].ParseFromString(
        nms_model._spec.description.output[1].SerializeToString()
    )

    pipeline.spec.specificationVersion = 5

    # Save
    f = str(output_path.with_suffix(".mlmodel"))
    pipeline_model = ct.models.MLModel(pipeline.spec)
    pipeline_model.input_description["image"] = "Input image"
    pipeline_model.input_description["iouThreshold"] = (
        f"(optional) IOU threshold override (default: {nms.iouThreshold})"
    )
    pipeline_model.input_description["confidenceThreshold"] = (
        f"(optional) Confidence threshold override (default: {nms.confidenceThreshold})"
    )
    pipeline_model.output_description["confidence"] = (
        "Boxes x Class confidence (see user-defined metadata 'classes')"
    )
    pipeline_model.output_description["coordinates"] = (
        "Boxes x [x, y, width, height] (relative to image size)"
    )
    pipeline_model.save(f)
    logger.info("CoreML pipeline: saved %s", f)
    return f


# --------------------------------------------------------------------------- #
# Exporter
# --------------------------------------------------------------------------- #

class CoreMLExporter(BaseExporter):
    """Export a PyTorch model to Apple CoreML ``.mlpackage`` format.

    The model is traced via ``torch.jit.trace`` and converted with
    ``coremltools.convert`` using an ``ImageType`` input.  Quantisation
    (INT8 k-means LUT or FP16 linear) is applied on macOS only.

    When NMS post-processing is requested the model is first wrapped in
    :class:`iOSModel` and, after conversion, an NMS pipeline is attached
    and saved as a ``.mlmodel``.
    """

    @property
    def format_name(self) -> str:
        return "coreml"

    @property
    def suffix(self) -> str:
        return ".mlpackage"

    def export(
        self,
        model: torch.nn.Module,
        sample_input: torch.Tensor,
        output_dir: str | Path,
        file_stem: str = "model",
        **kwargs: Any,
    ) -> str:
        """Export a PyTorch model to CoreML.

        Args:
            model: PyTorch model in eval mode.
            sample_input: Sample input tensor (BCHW).
            output_dir: Directory to write the exported file into.
            file_stem: Base filename without extension.
            **kwargs:
                int8 (bool): Quantise weights to 8-bit (k-means LUT).
                    Only effective on macOS.  Defaults to ``False``.
                half (bool): Quantise weights to FP16 (linear).  Only
                    effective on macOS.  Defaults to ``False``.
                nms (bool): Wrap the model with iOS-compatible NMS
                    post-processing.  Defaults to ``False``.
                names (dict[int, str] | None): Class names (required when
                    ``nms=True`` for the NMS pipeline labels).
                y: Raw PyTorch model output for shape inference on non-macOS
                    platforms when building the NMS pipeline.

        Returns:
            Absolute path to the exported ``.mlpackage`` (or ``.mlmodel``
            when an NMS pipeline is created).
        """
        import torch

        try:
            import coremltools as ct
        except ImportError as exc:
            raise ImportError(
                "CoreML export requires 'coremltools'. "
                "Install it with:  pip install coremltools"
            ) from exc

        int8: bool = kwargs.get("int8", False)
        half: bool = kwargs.get("half", False)
        nms: bool = kwargs.get("nms", False)
        names: dict | None = kwargs.get("names", None)
        y = kwargs.get("y", None)

        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        output_path = output_dir / f"{file_stem}{self.suffix}"

        logger.info(
            "CoreML export: converting model with coremltools %s ...",
            ct.__version__,
        )

        # Optionally wrap model for iOS NMS
        export_model = _make_ios_model(model, sample_input) if nms else model

        # Trace and convert
        ts = torch.jit.trace(export_model, sample_input, strict=False)
        ct_model = ct.convert(
            ts,
            inputs=[
                ct.ImageType(
                    "image",
                    shape=sample_input.shape,
                    scale=1 / 255.0,
                    bias=[0, 0, 0],
                )
            ],
        )

        # ------------------------------------------------------------------
        # Quantisation (macOS only)
        # ------------------------------------------------------------------
        bits, mode = (8, "kmeans_lut") if int8 else (16, "linear") if half else (32, None)
        if bits < 32:
            if MACOS:
                with warnings.catch_warnings():
                    warnings.filterwarnings("ignore", category=DeprecationWarning)
                    ct_model = (
                        ct.models.neural_network.quantization_utils.quantize_weights(
                            ct_model, bits, mode,
                        )
                    )
                logger.info("CoreML export: quantised to %d-bit (%s).", bits, mode)
            else:
                logger.warning(
                    "CoreML export: quantisation is only supported on macOS -- skipping."
                )

        ct_model.save(str(output_path))
        logger.info("CoreML export: saved %s", output_path)

        # ------------------------------------------------------------------
        # Optional NMS pipeline (.mlmodel)
        # ------------------------------------------------------------------
        if nms:
            if names is None:
                logger.warning(
                    "CoreML export: nms=True but no 'names' provided; "
                    "NMS pipeline will lack class labels."
                )
            if y is None:
                logger.warning(
                    "CoreML export: nms=True but no 'y' (model output) "
                    "provided for shape inference."
                )
            pipeline_path = _build_nms_pipeline(
                ct_model, sample_input, names, y, output_path,
            )
            return pipeline_path

        return str(output_path)
