"""Apple Silicon optimised CoreML export.

Provides a streamlined export path targeting the Apple Neural Engine (ANE)
with FP16 precision by default.  Unlike the full
:class:`~matrice_export.formats.coreml.CoreMLExporter` (which handles NMS
pipelines, image-type inputs, and platform detection), this module focuses
on *optimised inference* on Apple Silicon:

* ``ct.ComputeUnit.ALL`` enables CPU + GPU + ANE scheduling.
* ``ct.precision.FLOAT16`` reduces memory bandwidth and leverages ANE's
  native FP16 data-path.
* A lightweight validation helper compares the CoreML output against a
  PyTorch CPU baseline.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, Union

logger = logging.getLogger(__name__)


# --------------------------------------------------------------------------- #
# Export
# --------------------------------------------------------------------------- #

def export_coreml_optimized(
    model: Any,
    sample_input: Any,
    output_path: Union[str, Path],
    precision: str = "fp16",
) -> str:
    """Export a PyTorch model to CoreML optimised for Apple Neural Engine.

    The model is traced via ``torch.jit.trace`` and converted using
    ``coremltools.convert`` with compute-unit and precision settings tuned
    for Apple Silicon.

    Args:
        model: A ``torch.nn.Module`` in eval mode **or** a file-system path
            (``str`` / ``Path``) to a TorchScript model (``.pt`` / ``.pth``).
        sample_input: Sample input tensor (``torch.Tensor``) used for tracing
            and to determine the input shape for the CoreML model.
        output_path: Destination path for the exported ``.mlpackage``.  Parent
            directories are created automatically.
        precision: ``"fp16"`` (default) for half-precision or ``"fp32"`` for
            full single-precision.

    Returns:
        Absolute path to the saved ``.mlpackage`` as a string.

    Raises:
        ImportError: If ``coremltools`` or ``torch`` is not installed.
        ValueError: If *precision* is not ``"fp16"`` or ``"fp32"``.
        TypeError: If *model* is neither an ``nn.Module`` nor a valid path.
    """
    try:
        import coremltools as ct
    except ImportError as exc:
        raise ImportError(
            "CoreML optimised export requires 'coremltools'. "
            "Install it with:  pip install coremltools"
        ) from exc

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

    if precision not in ("fp16", "fp32"):
        raise ValueError(
            f"precision must be 'fp16' or 'fp32', got {precision!r}"
        )

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # ------------------------------------------------------------------ #
    # 1. Resolve model to a TorchScript trace
    # ------------------------------------------------------------------ #
    if isinstance(model, (str, Path)):
        logger.info("Loading TorchScript model from %s", model)
        traced = torch.jit.load(str(model), map_location="cpu")
    elif isinstance(model, torch.nn.Module):
        model.eval()
        logger.info("Tracing nn.Module for CoreML conversion ...")
        traced = torch.jit.trace(model.cpu(), sample_input.cpu(), strict=False)
    else:
        raise TypeError(
            f"model must be a torch.nn.Module or a path to a TorchScript "
            f"model, got {type(model).__name__}"
        )

    # ------------------------------------------------------------------ #
    # 2. Convert to CoreML
    # ------------------------------------------------------------------ #
    compute_precision = (
        ct.precision.FLOAT16 if precision == "fp16" else ct.precision.FLOAT32
    )

    logger.info(
        "Converting to CoreML (precision=%s, compute_units=ALL) ...",
        precision,
    )

    ct_model = ct.convert(
        traced,
        inputs=[ct.TensorType(shape=sample_input.shape)],
        compute_precision=compute_precision,
        compute_units=ct.ComputeUnit.ALL,  # CPU + GPU + ANE
    )

    # ------------------------------------------------------------------ #
    # 3. Save
    # ------------------------------------------------------------------ #
    save_path = str(output_path)
    ct_model.save(save_path)
    logger.info("CoreML optimised model saved to %s", save_path)
    return save_path


# --------------------------------------------------------------------------- #
# Validation
# --------------------------------------------------------------------------- #

def validate_coreml(
    model_path: Union[str, Path],
    sample_input: Any,
    pytorch_baseline: Any,
    atol: float = 1e-3,
    rtol: float = 1e-3,
) -> dict[str, Any]:
    """Validate a CoreML model against a PyTorch CPU baseline.

    Loads the ``.mlpackage``, runs a prediction with the given
    *sample_input*, and compares the result to *pytorch_baseline* using
    ``numpy.allclose``.

    Args:
        model_path: Path to the ``.mlpackage`` directory.
        sample_input: Input tensor (``torch.Tensor`` or ``numpy.ndarray``).
            Must match the shape expected by the CoreML model.
        pytorch_baseline: Expected output (``torch.Tensor`` or
            ``numpy.ndarray``) from the PyTorch model.
        atol: Absolute tolerance for ``numpy.allclose``.
        rtol: Relative tolerance for ``numpy.allclose``.

    Returns:
        A dict with keys:

        * ``status`` -- ``"pass"`` or ``"fail"``
        * ``max_abs_diff`` -- maximum element-wise absolute difference
        * ``mean_abs_diff`` -- mean element-wise absolute difference
        * ``atol`` -- the tolerance used
        * ``rtol`` -- the tolerance used

    Raises:
        ImportError: If ``coremltools`` or ``numpy`` is not installed.
        RuntimeError: If the CoreML prediction fails (e.g. running on
            Linux / Windows where the CoreML runtime is unavailable).
    """
    try:
        import coremltools as ct
    except ImportError as exc:
        raise ImportError(
            "CoreML validation requires 'coremltools'. "
            "Install it with:  pip install coremltools"
        ) from exc

    try:
        import numpy as np
    except ImportError as exc:
        raise ImportError(
            "CoreML validation requires 'numpy'. "
            "Install it with:  pip install numpy"
        ) from exc

    # Coerce inputs to numpy
    if hasattr(sample_input, "detach"):
        sample_np = sample_input.detach().cpu().numpy()
    else:
        sample_np = np.asarray(sample_input)

    if hasattr(pytorch_baseline, "detach"):
        baseline_np = pytorch_baseline.detach().cpu().numpy()
    else:
        baseline_np = np.asarray(pytorch_baseline)

    # Load and predict
    logger.info("Loading CoreML model from %s for validation ...", model_path)
    ct_model = ct.models.MLModel(str(model_path))

    input_name = ct_model.get_spec().description.input[0].name
    prediction = ct_model.predict({input_name: sample_np})

    # Extract first output
    output_name = ct_model.get_spec().description.output[0].name
    coreml_output = np.asarray(prediction[output_name])

    # Compare
    max_abs_diff = float(np.max(np.abs(coreml_output - baseline_np)))
    mean_abs_diff = float(np.mean(np.abs(coreml_output - baseline_np)))
    passed = bool(np.allclose(coreml_output, baseline_np, atol=atol, rtol=rtol))

    result: dict[str, Any] = {
        "status": "pass" if passed else "fail",
        "max_abs_diff": max_abs_diff,
        "mean_abs_diff": mean_abs_diff,
        "atol": atol,
        "rtol": rtol,
    }

    if passed:
        logger.info("CoreML validation passed (max_diff=%.6f)", max_abs_diff)
    else:
        logger.warning(
            "CoreML validation FAILED (max_diff=%.6f, atol=%.4f)",
            max_abs_diff,
            atol,
        )

    return result
