"""Convert SHARP PyTorch model to Core ML . mlmodel format. 

This script converts the SHARP (Sharp Monocular View Synthesis) model
from PyTorch (.pt) to Core ML (.mlmodel) format for deployment on Apple devices. 
"""

from __future__ import annotations

import argparse
import logging
from pathlib import Path
from typing import Any

import coremltools as ct
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

try:
    import onnx
    import onnxruntime as ort
    ONNX_AVAILABLE = True
except ImportError:
    ONNX_AVAILABLE = False
    print("Warning: ONNX not available, will use direct tracing")

# Import SHARP model components
from sharp. models import PredictorParams, create_predictor
from sharp.models.predictor import RGBGaussianPredictor

LOGGER = logging.getLogger(__name__)

DEFAULT_MODEL_URL = "https://ml-site.cdn-apple.com/models/sharp/sharp_2572gikvuh.pt"


class SharpModelWrapper(nn.Module):
    """Wrapper around RGBGaussianPredictor for Core ML export. 
    
    This wrapper simplifies the model interface for Core ML conversion by: 
    1.  Removing optional depth input (inference mode only)
    2. Flattening the Gaussians3D NamedTuple output to individual tensors
    3. Handling internal preprocessing
    """

    def __init__(self, predictor:  RGBGaussianPredictor):
        """Initialize the wrapper. 
        
        Args:
            predictor: The SHARP RGBGaussianPredictor model. 
        """
        super().__init__()
        self.predictor = predictor
        # Remove depth alignment for inference-only export
        # (depth alignment requires ground truth depth which is not available at inference)
        self.predictor.depth_alignment. scale_map_estimator = None
        
    def forward(
        self, 
        image: torch.Tensor, 
        disparity_factor: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: 
        """Run inference and return flattened Gaussian parameters.
        
        Args:
            image: Input image tensor of shape (1, 3, H, W) in range [0, 1].
            disparity_factor:  Disparity factor tensor of shape (1,).
            
        Returns:
            Tuple of tensors representing 3D Gaussians: 
                - mean_vectors: (1, N, 3) - 3D positions
                - singular_values: (1, N, 3) - scales
                - quaternions: (1, N, 4) - rotations
                - colors: (1, N, 3) - RGB colors
                - opacities: (1, N) - opacity values
        """
        # Run the predictor (depth=None for inference)
        gaussians = self.predictor(image, disparity_factor, depth=None)
        
        # Return as individual tensors (Core ML doesn't support NamedTuple)
        return (
            gaussians.mean_vectors,
            gaussians.singular_values,
            gaussians.quaternions,
            gaussians.colors,
            gaussians.opacities,
        )


class SafeClamp(nn.Module):
    """Safe clamp operation that avoids tracing issues."""

    def forward(self, x, min_val=1e-4, max_val=1e4):
        return torch.clamp(x, min=min_val, max=max_val)


class SafeDivision(nn.Module):
    """Safe division that avoids division by zero."""

    def forward(self, numerator, denominator):
        return numerator / torch.clamp(denominator, min=1e-8)


class SharpModelTraceable(nn.Module):
    """Fully traceable version of SHARP for Core ML conversion.

    This version removes all dynamic control flow and makes the model
    fully traceable with torch.jit.trace.
    """

    def __init__(self, predictor: RGBGaussianPredictor):
        """Initialize the traceable wrapper.

        Args:
            predictor: The SHARP RGBGaussianPredictor model.
        """
        super().__init__()
        # Copy all submodules
        self.init_model = predictor.init_model
        self. feature_model = predictor.feature_model
        self.monodepth_model = predictor.monodepth_model
        self.prediction_head = predictor.prediction_head
        self.gaussian_composer = predictor.gaussian_composer
        self.depth_alignment = predictor.depth_alignment

        # Replace problematic operations with custom modules
        self.safe_clamp = SafeClamp()
        self.safe_div = SafeDivision()
        
    def forward(
        self, 
        image: torch. Tensor, 
        disparity_factor: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Run inference with traceable forward pass.
        
        Args:
            image: Input image tensor of shape (1, 3, H, W) in range [0, 1]. 
            disparity_factor: Disparity factor tensor of shape (1,).
            
        Returns: 
            Tuple of 5 tensors representing 3D Gaussians. 
        """
        # Estimate depth using monodepth
        monodepth_output = self.monodepth_model(image)
        monodepth_disparity = monodepth_output. disparity

        # Convert disparity to depth with higher precision
        # Use tighter clamp bounds and higher precision intermediate computation
        disparity_factor_expanded = disparity_factor[: , None, None, None]
        
        # Cast to float64 for more precise division, then back to float32
        disparity_clamped = monodepth_disparity.clamp(min=1e-6, max=1e4)
        monodepth = disparity_factor_expanded.double() / disparity_clamped.double()
        monodepth = monodepth.float()

        # Apply depth alignment (inference mode)
        monodepth, _ = self.depth_alignment(monodepth, None, monodepth_output.decoder_features)

        # Initialize gaussians
        init_output = self.init_model(image, monodepth)
        
        # Extract features
        image_features = self. feature_model(
            init_output.feature_input, 
            encodings=monodepth_output.output_features
        )
        
        # Predict deltas
        delta_values = self.prediction_head(image_features)
        
        # Compose final gaussians
        gaussians = self.gaussian_composer(
            delta=delta_values,
            base_values=init_output. gaussian_base_values,
            global_scale=init_output.global_scale,
        )
        
        # Normalize quaternions for consistent validation and inference
        # This is critical for CoreML conversion accuracy
        quaternions = gaussians.quaternions
        
        # Use double precision for quaternion normalization to reduce numerical errors
        quaternions_fp64 = quaternions.double()
        quat_norm_sq = torch.sum(quaternions_fp64 * quaternions_fp64, dim=-1, keepdim=True)
        quat_norm = torch.sqrt(torch.clamp(quat_norm_sq, min=1e-16))
        quaternions_normalized = quaternions_fp64 / quat_norm
        
        # Apply sign canonicalization for consistent representation
        # Find the component with the largest absolute value
        abs_quat = torch.abs(quaternions_normalized)
        max_idx = torch.argmax(abs_quat, dim=-1, keepdim=True)
        
        # Create one-hot selector for the max component
        one_hot = torch.zeros_like(quaternions_normalized)
        one_hot.scatter_(-1, max_idx, 1.0)
        
        # Get the sign of the max component
        max_component_sign = torch.sum(quaternions_normalized * one_hot, dim=-1, keepdim=True)
        
        # Flip if negative (multiply by sign to canonicalize)
        sign = torch.sign(max_component_sign)
        sign = torch.where(sign == 0.0, torch.ones_like(sign), sign)  # Handle zero case
        quaternions = (quaternions_normalized * sign).float()
        
        return (
            gaussians.mean_vectors,
            gaussians.singular_values,
            quaternions,
            gaussians.colors,
            gaussians. opacities,
        )


def load_sharp_model(checkpoint_path: Path | None = None) -> RGBGaussianPredictor: 
    """Load SHARP model from checkpoint. 
    
    Args:
        checkpoint_path: Path to the . pt checkpoint file. 
                        If None, downloads the default model.
    
    Returns: 
        The loaded RGBGaussianPredictor model in eval mode.
    """
    if checkpoint_path is None: 
        LOGGER.info("Downloading default model from %s", DEFAULT_MODEL_URL)
        state_dict = torch.hub.load_state_dict_from_url(DEFAULT_MODEL_URL, progress=True)
    else:
        LOGGER.info("Loading checkpoint from %s", checkpoint_path)
        state_dict = torch.load(checkpoint_path, weights_only=True, map_location="cpu")
    
    # Create model with default parameters
    predictor = create_predictor(PredictorParams())
    predictor. load_state_dict(state_dict)
    predictor. eval()
    
    return predictor


def convert_via_onnx(
    predictor: RGBGaussianPredictor,
    output_path: Path,
    input_shape: tuple[int, int] = (1536, 1536),
    compute_precision: ct.precision = ct.precision.FLOAT32,
    compute_units: ct.ComputeUnit = ct.ComputeUnit.ALL,
    minimum_deployment_target: ct.target | None = None,
) -> ct.models.MLModel:
    """Convert SHARP model via ONNX intermediate format for better tracing.

    Args:
        predictor: The SHARP RGBGaussianPredictor model.
        output_path: Path to save the .mlmodel file.
        input_shape: Input image shape (height, width).
        compute_precision: Precision for compute (FLOAT16 or FLOAT32).
        compute_units: Target compute units.
        minimum_deployment_target: Minimum iOS/macOS deployment target.

    Returns:
        The converted Core ML model.
    """
    if not ONNX_AVAILABLE:
        raise ImportError("ONNX libraries not available. Install with: pip install onnx onnxruntime")

    LOGGER.info("Attempting ONNX-based conversion for better accuracy...")

    # Ensure depth alignment is disabled for inference
    predictor.depth_alignment.scale_map_estimator = None

    # Create traceable wrapper
    model_wrapper = SharpModelTraceable(predictor)
    model_wrapper.eval()

    # Pre-warm the model
    LOGGER.info("Pre-warming model...")
    with torch.no_grad():
        for _ in range(3):
            warm_image = torch.randn(1, 3, input_shape[0], input_shape[1])
            warm_disparity = torch.tensor([1.0])
            _ = model_wrapper(warm_image, warm_disparity)

    # Create example inputs
    height, width = input_shape
    torch.manual_seed(42)
    example_image = torch.randn(1, 3, height, width)
    example_disparity_factor = torch.tensor([1.0])

    # Export to ONNX with external data to handle large models
    onnx_path = output_path.with_suffix('.onnx')
    LOGGER.info(f"Exporting to ONNX: {onnx_path}")

    try:
        # Export with external data format to handle large models (>2GB)
        torch.onnx.export(
            model_wrapper,
            (example_image, example_disparity_factor),
            str(onnx_path),  # Must be string path for external data format
            export_params=True,
            verbose=False,
            input_names=['image', 'disparity_factor'],
            output_names=[
                'mean_vectors_3d_positions',
                'singular_values_scales',
                'quaternions_rotations',
                'colors_rgb_linear',
                'opacities_alpha_channel'
            ],
            dynamic_axes={
                'mean_vectors_3d_positions': {1: 'num_gaussians'},
                'singular_values_scales': {1: 'num_gaussians'},
                'quaternions_rotations': {1: 'num_gaussians'},
                'colors_rgb_linear': {1: 'num_gaussians'},
                'opacities_alpha_channel': {1: 'num_gaussians'}
            },
            opset_version=17,  # Use latest stable opset
        )
        
        # For models >2GB, save with external data format
        try:
            import onnx
            model_proto = onnx.load(str(onnx_path))
            # Check if model is larger than 2GB (need external data)
            model_size = model_proto.ByteSize()
            if model_size > 2e9:  # 2GB
                LOGGER.info(f"Model size {model_size/1e9:.2f}GB > 2GB, converting to external data format...")
                onnx.save_model(
                    model_proto,
                    str(onnx_path),
                    save_as_external_data=True,
                    all_tensors_to_one_file=True,
                    location=f"{onnx_path.stem}.onnx.data",
                    size_threshold=1024,
                    convert_attribute=False,
                )
                LOGGER.info("Successfully saved with external data format")
        except Exception as e:
            LOGGER.warning(f"Could not check/convert to external data format: {e}")
            
        LOGGER.info("ONNX export successful")
    except Exception as e:
        LOGGER.error(f"ONNX export failed: {e}")
        raise

    # Verify ONNX model (skip for large models >2GB)
    try:
        # For large models, check without loading external data
        onnx.checker.check_model(str(onnx_path))
        LOGGER.info("ONNX model validation passed")
    except Exception as e:
        LOGGER.warning(f"ONNX model validation skipped for large model: {e}")

    # Convert ONNX to Core ML
    LOGGER.info("Converting ONNX to Core ML...")

    # Define input types for Core ML
    inputs = [
        ct.TensorType(
            name="image",
            shape=(1, 3, height, width),
            dtype=np.float32,
        ),
        ct.TensorType(
            name="disparity_factor",
            shape=(1,),
            dtype=np.float32,
        ),
    ]

    # Convert from ONNX - must specify source="pytorch" for proper conversion
    mlmodel = ct.convert(
        str(onnx_path),
        source="pytorch",  # Specify PyTorch as source framework
        inputs=inputs,
        convert_to="mlprogram",
        compute_precision=compute_precision,
        compute_units=compute_units,
        minimum_deployment_target=minimum_deployment_target,
    )

    # Add metadata
    mlmodel.author = "Apple Inc. (via ONNX)"
    mlmodel.license = "See LICENSE_MODEL in ml-sharp repository"
    mlmodel.short_description = (
        "SHARP: Sharp Monocular View Synthesis - Predicts 3D Gaussian splats from a single image (ONNX conversion)"
    )
    mlmodel.version = "1.0.0"

    # Add output descriptions
    spec = mlmodel.get_spec()
    output_descriptions = {
        "mean_vectors_3d_positions": (
            "3D positions of Gaussian splats in normalized device coordinates (NDC). "
            "Shape: (1, N, 3), where N is the number of Gaussians."
        ),
        "singular_values_scales": (
            "Scale factors for each Gaussian along its principal axes. "
            "Represents size and anisotropy. Shape: (1, N, 3)."
        ),
        "quaternions_rotations": (
            "Rotation of each Gaussian as a unit quaternion [w, x, y, z]. "
            "Used to orient the ellipsoid. Shape: (1, N, 4)."
        ),
        "colors_rgb_linear": (
            "RGB color values in linear RGB space (not gamma-corrected). "
            "Shape: (1, N, 3), with range [0, 1]."
        ),
        "opacities_alpha_channel": (
            "Opacity value per Gaussian (alpha channel), used for blending. "
            "Shape: (1, N), where values are in [0, 1]."
        ),
    }

    # Update output names and descriptions
    for i, name in enumerate([
        'mean_vectors_3d_positions',
        'singular_values_scales',
        'quaternions_rotations',
        'colors_rgb_linear',
        'opacities_alpha_channel'
    ]):
        if i < len(spec.description.output):
            output = spec.description.output[i]
            output.name = name
            output.shortDescription = output_descriptions[name]

    # Validate output names are set correctly
    LOGGER.info("Output names after update: %s", [o.name for o in spec.description.output])

    # Save the model
    LOGGER.info(f"Saving Core ML model to {output_path}")
    mlmodel.save(str(output_path))

    # Clean up ONNX file
    try:
        onnx_path.unlink()
        LOGGER.info("Cleaned up temporary ONNX file")
    except Exception as e:
        LOGGER.warning(f"Could not clean up ONNX file: {e}")

    return mlmodel


def convert_to_coreml(
    predictor: RGBGaussianPredictor,
    output_path: Path,
    input_shape: tuple[int, int] = (1536, 1536),
    compute_precision: ct.precision = ct.precision.FLOAT16,
    compute_units: ct.ComputeUnit = ct.ComputeUnit.ALL,
    minimum_deployment_target: ct.target | None = None,
) -> ct.models.MLModel:
    """Convert SHARP model to Core ML format.
    
    Args:
        predictor: The SHARP RGBGaussianPredictor model.
        output_path: Path to save the . mlmodel file. 
        input_shape:  Input image shape (height, width). Default is (1536, 1536).
        compute_precision:  Precision for compute (FLOAT16 or FLOAT32).
        compute_units: Target compute units (ALL, CPU_AND_GPU, CPU_ONLY, etc.).
        minimum_deployment_target:  Minimum iOS/macOS deployment target.
        
    Returns: 
        The converted Core ML model.
    """
    LOGGER.info("Preparing model for Core ML conversion...")

    # Ensure depth alignment is disabled for inference
    predictor.depth_alignment.scale_map_estimator = None

    # Create traceable wrapper
    model_wrapper = SharpModelTraceable(predictor)
    model_wrapper.eval()

    # Pre-warm the model with a few forward passes for better tracing
    LOGGER.info("Pre-warming model for better tracing...")
    with torch.no_grad():
        for _ in range(3):
            warm_image = torch.randn(1, 3, input_shape[0], input_shape[1])
            warm_disparity = torch.tensor([1.0])
            _ = model_wrapper(warm_image, warm_disparity)

    # Create deterministic example inputs for tracing (same as validation)
    height, width = input_shape
    torch.manual_seed(42)  # Use same seed as validation for consistency
    example_image = torch.randn(1, 3, height, width)
    example_disparity_factor = torch.tensor([1.0])
    
    LOGGER.info("Attempting torch.jit.script for better tracing...")
    try:
        with torch.no_grad():
            scripted_model = torch.jit.script(model_wrapper)
        LOGGER.info("torch.jit.script succeeded, using scripted model")
        traced_model = scripted_model
    except Exception as e:
        LOGGER.warning(f"torch.jit.script failed: {e}")
        LOGGER.info("Falling back to torch.jit.trace...")
        with torch.no_grad():
            traced_model = torch.jit.trace(
                model_wrapper,
                (example_image, example_disparity_factor),
                strict=False,  # Allow some flexibility for complex models
                check_trace=False,  # Skip trace checking to allow more flexibility
            )
    
    LOGGER.info("Converting traced model to Core ML...")
    
    # Define input types for Core ML
    inputs = [
        ct.TensorType(
            name="image",
            shape=(1, 3, height, width),
            dtype=np.float32,
        ),
        ct.TensorType(
            name="disparity_factor",
            shape=(1,),
            dtype=np.float32,
        ),
    ]
    
    # Define output names with clear, descriptive labels
    output_names = [
        "mean_vectors_3d_positions",         # 3D positions (NDC space)
        "singular_values_scales",            # Scale parameters (diagonal of covariance)
        "quaternions_rotations",             # Rotation as quaternions
        "colors_rgb_linear",                 # RGB colors in linear color space
        "opacities_alpha_channel",           # Opacity values (alpha)
    ]

    # Define outputs with proper names for Core ML conversion
    outputs = [
        ct.TensorType(name=output_names[0], dtype=np.float32),
        ct.TensorType(name=output_names[1], dtype=np.float32),
        ct.TensorType(name=output_names[2], dtype=np.float32),
        ct.TensorType(name=output_names[3], dtype=np.float32),
        ct.TensorType(name=output_names[4], dtype=np.float32),
    ]

    # Set up conversion config
    conversion_kwargs:  dict[str, Any] = {
        "inputs": inputs,
        "outputs": outputs,  # Specify output names during conversion
        "convert_to":  "mlprogram",  # Use ML Program format for better performance
        "compute_precision": compute_precision,
        "compute_units": compute_units,
    }

    if minimum_deployment_target is not None:
        conversion_kwargs["minimum_deployment_target"] = minimum_deployment_target

    # Convert to Core ML
    mlmodel = ct.convert(
        traced_model,
        **conversion_kwargs,
    )
    
    # Add metadata
    mlmodel.author = "Apple Inc."
    mlmodel.license = "See LICENSE_MODEL in ml-sharp repository"
    mlmodel.short_description = (
        "SHARP: Sharp Monocular View Synthesis - Predicts 3D Gaussian splats from a single image"
    )
    mlmodel.version = "1.0.0"

    # Update output names and descriptions via spec BEFORE saving
    spec = mlmodel.get_spec()

    # Input descriptions
    input_descriptions = {
        "image": "RGB image normalized to [0, 1], shape (1, 3, H, W)",
        "disparity_factor": "Focal length / image width ratio, shape (1,)",
    }

    # Output descriptions with clear intent and units
    output_descriptions = {
        "mean_vectors_3d_positions": (
            "3D positions of Gaussian splats in normalized device coordinates (NDC). "
            "Shape: (1, N, 3), where N is the number of Gaussians."
        ),
        "singular_values_scales": (
            "Scale factors for each Gaussian along its principal axes. "
            "Represents size and anisotropy. Shape: (1, N, 3)."
        ),
        "quaternions_rotations": (
            "Rotation of each Gaussian as a unit quaternion [w, x, y, z]. "
            "Used to orient the ellipsoid. Shape: (1, N, 4)."
        ),
        "colors_rgb_linear": (
            "RGB color values in linear RGB space (not gamma-corrected). "
            "Shape: (1, N, 3), with range [0, 1]."
        ),
        "opacities_alpha_channel": (
            "Opacity value per Gaussian (alpha channel), used for blending. "
            "Shape: (1, N), where values are in [0, 1]."
        ),
    }

    # Update output names and descriptions
    for i, name in enumerate(output_names):
        if i < len(spec.description.output):
            output = spec.description.output[i]
            output.name = name  # Update name
            output.shortDescription = output_descriptions[name]  # Add description

    # Validate output names are set correctly
    LOGGER.info("Output names after update: %s", [o.name for o in spec.description.output])

    # Save the model with correct names
    LOGGER.info("Saving Core ML model to %s", output_path)
    mlmodel.save(str(output_path))

    return mlmodel


def convert_to_coreml_with_preprocessing(
    predictor: RGBGaussianPredictor,
    output_path: Path,
    input_shape:  tuple[int, int] = (1536, 1536),
) -> ct.models. MLModel:
    """Convert SHARP model to Core ML with built-in image preprocessing.
    
    This version includes image normalization as part of the model,
    accepting uint8 images as input.
    
    Args:
        predictor: The SHARP RGBGaussianPredictor model.
        output_path: Path to save the .mlmodel file.
        input_shape: Input image shape (height, width).
        
    Returns: 
        The converted Core ML model.
    """
    
    class SharpWithPreprocessing(nn.Module):
        """SHARP model with integrated preprocessing."""
        
        def __init__(self, base_model: SharpModelTraceable):
            super().__init__()
            self.base_model = base_model
            
        def forward(
            self, 
            image: torch. Tensor,
            disparity_factor: torch. Tensor
        ) -> tuple[torch. Tensor, torch. Tensor, torch. Tensor, torch. Tensor, torch. Tensor]:
            # Normalize image from [0, 255] to [0, 1]
            image_normalized = image / 255.0
            return self.base_model(image_normalized, disparity_factor)
    
    model_wrapper = SharpWithPreprocessing(SharpModelTraceable(predictor))
    model_wrapper.eval()
    
    height, width = input_shape
    example_image = torch.randint(0, 256, (1, 3, height, width), dtype=torch.float32)
    example_disparity_factor = torch.tensor([1.0])
    
    LOGGER.info("Tracing model with preprocessing...")
    with torch.no_grad():
        traced_model = torch.jit.trace(
            model_wrapper,
            (example_image, example_disparity_factor),
            strict=False,
        )
    
    inputs = [
        ct.ImageType(
            name="image",
            shape=(1, 3, height, width),
            scale=1.0,  # Will be normalized in the model
            color_layout=ct.colorlayout.RGB,
        ),
        ct.TensorType(
            name="disparity_factor",
            shape=(1,),
            dtype=np. float32,
        ),
    ]

    # Define output names with clear, descriptive labels
    output_names = [
        "mean_vectors_3d_positions",         # 3D positions (NDC space)
        "singular_values_scales",            # Scale parameters (diagonal of covariance)
        "quaternions_rotations",             # Rotation as quaternions
        "colors_rgb_linear",                 # RGB colors in linear color space
        "opacities_alpha_channel",           # Opacity values (alpha)
    ]

    # Define outputs with proper names for Core ML conversion
    outputs = [
        ct.TensorType(name=output_names[0], dtype=np.float32),
        ct.TensorType(name=output_names[1], dtype=np.float32),
        ct.TensorType(name=output_names[2], dtype=np.float32),
        ct.TensorType(name=output_names[3], dtype=np.float32),
        ct.TensorType(name=output_names[4], dtype=np.float32),
    ]

    mlmodel = ct.convert(
        traced_model,
        inputs=inputs,
        outputs=outputs,  # Specify output names during conversion
        convert_to="mlprogram",
        compute_precision=ct.precision.FLOAT16,
    )

    mlmodel.author = "Apple Inc."
    mlmodel.short_description = "SHARP model with integrated image preprocessing"
    mlmodel.version = "1.0.0"

    # Define output names with clear, descriptive labels
    output_names = [
        "mean_vectors_3d_positions",         # 3D positions (NDC space)
        "singular_values_scales",            # Scale parameters (diagonal of covariance)
        "quaternions_rotations",             # Rotation as quaternions
        "colors_rgb_linear",                 # RGB colors in linear color space
        "opacities_alpha_channel",           # Opacity values (alpha)
    ]

    # Output descriptions with clear intent and units
    output_descriptions = {
        "mean_vectors_3d_positions": (
            "3D positions of Gaussian splats in normalized device coordinates (NDC). "
            "Shape: (1, N, 3), where N is the number of Gaussians."
        ),
        "singular_values_scales": (
            "Scale factors for each Gaussian along its principal axes. "
            "Represents size and anisotropy. Shape: (1, N, 3)."
        ),
        "quaternions_rotations": (
            "Rotation of each Gaussian as a unit quaternion [w, x, y, z]. "
            "Used to orient the ellipsoid. Shape: (1, N, 4)."
        ),
        "colors_rgb_linear": (
            "RGB color values in linear RGB space (not gamma-corrected). "
            "Shape: (1, N, 3), with range [0, 1]."
        ),
        "opacities_alpha_channel": (
            "Opacity value per Gaussian (alpha channel), used for blending. "
            "Shape: (1, N), where values are in [0, 1]."
        ),
    }

    # Update output names and descriptions via spec BEFORE saving
    spec = mlmodel.get_spec()

    for i, name in enumerate(output_names):
        if i < len(spec.description.output):
            output = spec.description.output[i]
            output.name = name  # Update name
            output.shortDescription = output_descriptions[name]  # Add description

    # Validate output names are set correctly
    LOGGER.info("Output names after update: %s", [o.name for o in spec.description.output])

    # Save the model with correct names
    mlmodel.save(str(output_path))

    return mlmodel


def validate_coreml_model(
    mlmodel: ct.models.MLModel,
    pytorch_model: RGBGaussianPredictor,
    input_shape:  tuple[int, int] = (1536, 1536),
    tolerance: float = 0.01,  # Much tighter tolerance
) -> bool:
    """Validate Core ML model outputs against PyTorch model.
    
    Args:
        mlmodel: The Core ML model to validate.
        pytorch_model: The original PyTorch model. 
        input_shape: Input image shape (height, width).
        tolerance: Maximum allowed difference between outputs.
        
    Returns: 
        True if validation passes, False otherwise. 
    """
    LOGGER.info("Validating Core ML model against PyTorch...")
    
    height, width = input_shape
    
    # Create test input
    np. random.seed(42)
    test_image_np = np.random.rand(1, 3, height, width).astype(np.float32)
    test_disparity = np.array([1.0], dtype=np.float32)
    
    # Run PyTorch model
    test_image_pt = torch.from_numpy(test_image_np)
    test_disparity_pt = torch.from_numpy(test_disparity)
    
    traceable_wrapper = SharpModelTraceable(pytorch_model)
    traceable_wrapper.eval()
    
    with torch.no_grad():
        pt_outputs = traceable_wrapper(test_image_pt, test_disparity_pt)

    # Run Core ML model
    coreml_inputs = {
        "image":  test_image_np,
        "disparity_factor": test_disparity,
    }
    coreml_outputs = mlmodel.predict(coreml_inputs)

    # Debug: Print shapes and keys
    LOGGER.info(f"PyTorch outputs shapes: {[o.shape for o in pt_outputs]}")
    LOGGER.info(f"Core ML outputs keys: {list(coreml_outputs.keys())}")
    
    # Compare outputs with per-output tolerances
    output_names = ["mean_vectors_3d_positions", "singular_values_scales", "quaternions_rotations", "colors_rgb_linear", "opacities_alpha_channel"]
    
    # Define tighter tolerances per output type
    tolerances = {
        "mean_vectors_3d_positions": 0.001,     # 1mm precision
        "singular_values_scales": 0.0001,       # 0.01% scale precision  
        "quaternions_rotations": 2.0,           # Raw component diff (use angular for actual check)
        "colors_rgb_linear": 0.002,             # ~0.5/255 in 8-bit
        "opacities_alpha_channel": 0.005,       # ~1/255 in 8-bit
    }
    
    # Angular tolerances for quaternions (in degrees)
    angular_tolerances = {
        "mean": 0.01,    # Mean angular error should be < 0.01°
        "p99": 0.5,      # 99th percentile should be < 0.5°
        "max": 10.0,     # Max angular error should be < 10° (allows a few outliers)
    }
    
    all_passed = True
    
    # Additional diagnostics for depth/position analysis
    LOGGER.info("=== Depth/Position Statistics ===")
    pt_positions = pt_outputs[0].numpy()
    coreml_positions = coreml_outputs[list(coreml_outputs.keys())[0] if "mean_vectors" not in list(coreml_outputs.keys())[0] else [k for k in coreml_outputs.keys() if "mean_vectors" in k][0]]
    
    LOGGER.info(f"PyTorch positions - X range: [{pt_positions[..., 0].min():.4f}, {pt_positions[..., 0].max():.4f}], mean: {pt_positions[..., 0].mean():.4f}")
    LOGGER.info(f"PyTorch positions - Y range: [{pt_positions[..., 1].min():.4f}, {pt_positions[..., 1].max():.4f}], mean: {pt_positions[..., 1].mean():.4f}")
    LOGGER.info(f"PyTorch positions - Z range: [{pt_positions[..., 2].min():.4f}, {pt_positions[..., 2].max():.4f}], mean: {pt_positions[..., 2].mean():.4f}, std: {pt_positions[..., 2].std():.4f}")
    
    LOGGER.info(f"CoreML positions - X range: [{coreml_positions[..., 0].min():.4f}, {coreml_positions[..., 0].max():.4f}], mean: {coreml_positions[..., 0].mean():.4f}")
    LOGGER.info(f"CoreML positions - Y range: [{coreml_positions[..., 1].min():.4f}, {coreml_positions[..., 1].max():.4f}], mean: {coreml_positions[..., 1].mean():.4f}")
    LOGGER.info(f"CoreML positions - Z range: [{coreml_positions[..., 2].min():.4f}, {coreml_positions[..., 2].max():.4f}], mean: {coreml_positions[..., 2].mean():.4f}, std: {coreml_positions[..., 2].std():.4f}")
    
    z_diff = np.abs(pt_positions[..., 2] - coreml_positions[..., 2])
    LOGGER.info(f"Z-coordinate difference - max: {z_diff.max():.6f}, mean: {z_diff.mean():.6f}, std: {z_diff.std():.6f}")
    LOGGER.info("=================================")
    
    for i, name in enumerate(output_names):
        pt_output = pt_outputs[i]. numpy()
        
        # Find matching Core ML output (names may have suffixes)
        coreml_key = None
        for key in coreml_outputs:
            if name in key. lower() or f"var_{i}" in key: 
                coreml_key = key
                break
        
        if coreml_key is None: 
            # Try by index
            coreml_key = list(coreml_outputs.keys())[i]
        
        coreml_output = coreml_outputs[coreml_key]
        
        # Special handling for quaternions - account for sign ambiguity
        if name == "quaternions_rotations":
            # Normalize both quaternion outputs to ensure they're unit quaternions
            pt_quat_norm = np.linalg.norm(pt_output, axis=-1, keepdims=True)
            pt_output_normalized = pt_output / np.clip(pt_quat_norm, 1e-12, None)
            
            coreml_quat_norm = np.linalg.norm(coreml_output, axis=-1, keepdims=True)
            coreml_output_normalized = coreml_output / np.clip(coreml_quat_norm, 1e-12, None)
            
            # Canonicalize sign: handle edge cases where w ≈ 0
            # When w is near zero, use the largest magnitude component for canonicalization
            def canonicalize_quaternion(q):
                """Canonicalize quaternion to ensure unique representation."""
                # Find the component with the largest absolute value
                abs_q = np.abs(q)
                max_component_idx = np.argmax(abs_q, axis=-1, keepdims=True)
                
                # Create a selector for the max component
                selector = np.zeros_like(q)
                np.put_along_axis(selector, max_component_idx, 1, axis=-1)
                
                # Get the sign of the max component
                max_component_sign = np.sum(q * selector, axis=-1, keepdims=True)
                
                # Flip quaternion if the max component is negative
                return np.where(max_component_sign < 0, -q, q)
            
            pt_output_canonical = canonicalize_quaternion(pt_output_normalized)
            coreml_output_canonical = canonicalize_quaternion(coreml_output_normalized)
            
            # Now compute differences with canonicalized quaternions
            diff = np.abs(pt_output_canonical - coreml_output_canonical)
            max_diff = np.max(diff)
            mean_diff = np.mean(diff)
            
            # Also compute angular difference for better insight
            # Angular distance = 2 * arccos(|dot(q1, q2)|)
            dot_products = np.sum(pt_output_canonical * coreml_output_canonical, axis=-1)
            dot_products = np.clip(np.abs(dot_products), 0.0, 1.0)
            angular_diff_rad = 2 * np.arccos(dot_products)
            angular_diff_deg = np.degrees(angular_diff_rad)
            max_angular = np.max(angular_diff_deg)
            mean_angular = np.mean(angular_diff_deg)
            p99_angular = np.percentile(angular_diff_deg, 99)
            
            # Check angular tolerances (more meaningful for rotations)
            quat_passed = True
            failure_reasons = []
            
            if mean_angular > angular_tolerances["mean"]:
                quat_passed = False
                failure_reasons.append(f"mean angular {mean_angular:.4f}° > {angular_tolerances['mean']:.4f}°")
            
            if p99_angular > angular_tolerances["p99"]:
                quat_passed = False
                failure_reasons.append(f"p99 angular {p99_angular:.4f}° > {angular_tolerances['p99']:.4f}°")
            
            if max_angular > angular_tolerances["max"]:
                quat_passed = False
                failure_reasons.append(f"max angular {max_angular:.4f}° > {angular_tolerances['max']:.4f}°")
            
            if not quat_passed:
                LOGGER.warning(
                    "Output %s: max diff %.6f, mean diff %.6f, p99 diff %.6f, max angular diff %.4f°, mean angular diff %.4f°, p99 angular %.4f° (FAILED: %s)",
                    name, max_diff, mean_diff, np.percentile(diff, 99), max_angular, mean_angular, p99_angular, "; ".join(failure_reasons)
                )
                all_passed = False
            else:
                LOGGER.info(
                    "Output %s: max diff %.6f, mean diff %.6f, p99 diff %.6f, max angular diff %.4f°, mean angular diff %.4f°, p99 angular %.4f° (PASSED)",
                    name, max_diff, mean_diff, np.percentile(diff, 99), max_angular, mean_angular, p99_angular
                )
        else:
            diff = np.abs(pt_output - coreml_output)
            max_diff = np.max(diff)
            mean_diff = np.mean(diff)
            p99_diff = np.percentile(diff, 99)
            
            # Compute relative error for better insight
            relative_diff = diff / (np.abs(pt_output) + 1e-8)
            max_relative = np.max(relative_diff)
            mean_relative = np.mean(relative_diff)
            
            output_tolerance = tolerances.get(name, tolerance)
            
            if max_diff > output_tolerance:
                LOGGER.warning(
                    "Output %s: max diff %.6f, mean diff %.6f, p99 diff %.6f, max rel %.4f%%, mean rel %.4f%% (FAILED, tolerance=%.6f)",
                    name, max_diff, mean_diff, p99_diff, max_relative * 100, mean_relative * 100, output_tolerance
                )
                all_passed = False
            else:
                LOGGER.info(
                    "Output %s: max diff %.6f, mean diff %.6f, p99 diff %.6f, max rel %.4f%%, mean rel %.4f%% (PASSED)",
                    name, max_diff, mean_diff, p99_diff, max_relative * 100, mean_relative * 100
                )
    
    return all_passed


def main():
    """Main conversion script."""
    parser = argparse.ArgumentParser(
        description="Convert SHARP PyTorch model to Core ML format"
    )
    parser.add_argument(
        "-c", "--checkpoint",
        type=Path,
        default=None,
        help="Path to PyTorch checkpoint.  Downloads default if not provided.",
    )
    parser.add_argument(
        "-o", "--output",
        type=Path,
        default=Path("sharp.mlpackage"),
        help="Output path for Core ML model (default: sharp.mlpackage)",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=1536,
        help="Input image height (default:  1536)",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=1536,
        help="Input image width (default: 1536)",
    )
    parser.add_argument(
        "--precision",
        choices=["float16", "float32"],
        default="float32",
        help="Compute precision (default: float32)",
    )
    parser.add_argument(
        "--validate",
        action="store_true",
        help="Validate Core ML model against PyTorch",
    )
    parser.add_argument(
        "--with-preprocessing",
        action="store_true",
        help="Include image preprocessing (uint8 -> float normalization)",
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Enable verbose logging",
    )
    
    args = parser.parse_args()
    
    # Configure logging
    logging.basicConfig(
        level=logging. DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )
    
    # Load PyTorch model
    LOGGER.info("Loading SHARP model...")
    predictor = load_sharp_model(args. checkpoint)
    
    # Convert to Core ML
    input_shape = (args.height, args. width)
    precision = ct.precision.FLOAT16 if args. precision == "float16" else ct.precision.FLOAT32
    
    if args.with_preprocessing:
        LOGGER.info("Converting with integrated preprocessing...")
        mlmodel = convert_to_coreml_with_preprocessing(
            predictor,
            args.output,
            input_shape=input_shape,
        )
    else:
        # Try ONNX conversion first for better accuracy
        if ONNX_AVAILABLE:
            try:
                LOGGER.info("Attempting ONNX-based conversion...")
                mlmodel = convert_via_onnx(
                    predictor,
                    args.output,
                    input_shape=input_shape,
                    compute_precision=precision,
                )
                LOGGER.info("ONNX conversion successful!")
            except Exception as e:
                LOGGER.warning(f"ONNX conversion failed: {e}")
                LOGGER.info("Falling back to direct tracing...")
                mlmodel = convert_to_coreml(
                    predictor,
                    args.output,
                    input_shape=input_shape,
                    compute_precision=precision,
                )
        else:
            LOGGER.info("ONNX not available, using direct tracing...")
            mlmodel = convert_to_coreml(
                predictor,
                args.output,
                input_shape=input_shape,
                compute_precision=precision,
            )
    
    LOGGER.info("Core ML model saved to %s", args.output)
    
    # Validate if requested
    if args.validate:
        if validate_coreml_model(mlmodel, predictor, input_shape):
            LOGGER.info("✓ Validation passed!")
        else: 
            LOGGER.error("✗ Validation failed!")
            return 1
    
    LOGGER.info("Conversion complete!")
    return 0


if __name__ == "__main__":
    exit(main())