"""Optional fixed-tile Core ML decoder with the native MLX tile assembly.

Core ML uses FP16 internally. This is an explicitly approximate alternative to
the native BF16 decoder; use the saved-latent image comparisons before choosing.
"""

from __future__ import annotations

from pathlib import Path

import mlx.core as mx
import numpy as np

from .vae import VAE


class CoreMLVAE(VAE):
    @staticmethod
    def validate_metadata(meta, source_fingerprint, tile_batch):
        if meta.get("h3.decoder.format") != "h3-coreml-tile-v1":
            raise ValueError("This ML Package lacks H3 decoder provenance metadata")
        if meta.get("h3.decoder.source") != source_fingerprint:
            raise ValueError(
                "Core ML decoder does not match this source VAE checkpoint"
            )
        fixed_batch = int(meta["h3.decoder.batch"])
        if tile_batch != fixed_batch:
            raise ValueError(f"ML Package was exported for --vae-batch {fixed_batch}")
        return fixed_batch

    def __init__(
        self,
        path,
        *,
        package,
        compute_units="CPU_AND_GPU",
        tile_batch=8,
        dtype="bfloat16",
        compile_blocks=False,
    ):
        if dtype != "bfloat16":
            raise ValueError(
                "Core ML tile assembly currently requires the native BF16 policy"
            )
        import coremltools as ct

        super().__init__(path, tile_batch=tile_batch, dtype=dtype)
        self.package = Path(package)
        self.model = ct.models.MLModel(
            str(self.package),
            compute_units=getattr(ct.ComputeUnit, compute_units),
            optimization_hints={
                "specializationStrategy": ct.SpecializationStrategy.FastPrediction
            },
        )
        meta = self.model.user_defined_metadata
        self.fixed_batch = self.validate_metadata(
            meta, self.store.fingerprint(), tile_batch
        )

    def decode_tiles(self, raw):
        if raw.shape[1:] != (24, 16, 16) or raw.shape[0] > self.fixed_batch:
            raise ValueError(
                "Core ML export requires native 256px tiles and its fixed batch size"
            )
        count = raw.shape[0]
        # BF16 source values are exactly representable by FP16 in the ordinary
        # decoder range. Preserve the native input rounding before conversion.
        x = raw.astype(mx.bfloat16).astype(mx.float16)
        if count < self.fixed_batch:
            x = mx.pad(x, ((0, self.fixed_batch - count), (0, 0), (0, 0), (0, 0)))
        pixels = self.model.predict({"raw": np.ascontiguousarray(np.asarray(x))})[
            "pixels"
        ]
        if not np.isfinite(pixels).all():
            raise FloatingPointError("Core ML decoder produced non-finite pixels")
        return mx.array(np.ascontiguousarray(pixels[:count])).astype(mx.bfloat16)

    def clear(self):
        self.model = None
        super().clear()


def source_metadata(path, batch):
    """Metadata shared with the exporter, without importing the converter."""
    from .weights import WeightStore

    return {
        "h3.decoder.format": "h3-coreml-tile-v1",
        "h3.decoder.source": WeightStore(path).fingerprint(),
        "h3.decoder.batch": str(batch),
        "h3.decoder.precision": "Source BF16 weights and inputs, FP16 internal arithmetic, BF16 tile assembly",
    }