# Copyright 2025 Yakhyokhuja Valikhujaev
# Author: Yakhyokhuja Valikhujaev
# GitHub: https://github.com/yakhyo
#
# MLX Backend Support added by Claude Code

"""
MLX utilities for UniFace.

This module provides utilities for running inference on Apple Silicon using MLX.
MLX is Apple's machine learning framework optimized for M1/M2/M3/M4 chips with
unified memory architecture.
"""

import platform
from typing import Any, Dict, Optional

import numpy as np

from uniface.log import Logger

__all__ = [
    'is_mlx_available',
    'is_apple_silicon',
    'get_mlx_device',
    'load_mlx_weights',
    'load_mlx_fused_weights',
    'create_mlx_model',
    'to_mlx_array',
    'to_numpy',
    'synchronize',
]

# Lazy import for MLX
_mlx_available: Optional[bool] = None
_mx = None
_nn = None


def _import_mlx():
    """Lazily import MLX modules."""
    global _mx, _nn, _mlx_available
    if _mlx_available is None:
        try:
            import mlx.core as mx
            import mlx.nn as nn

            _mx = mx
            _nn = nn
            _mlx_available = True
            Logger.debug('MLX successfully imported')
        except ImportError:
            _mlx_available = False
            Logger.debug('MLX not available - falling back to ONNX')
    return _mlx_available


def is_apple_silicon() -> bool:
    """
    Check if running on Apple Silicon (M1/M2/M3/M4).

    Returns:
        bool: True if running on Apple Silicon, False otherwise.
    """
    if platform.system() != 'Darwin':
        return False

    # Check for ARM64 architecture (Apple Silicon)
    return platform.machine() == 'arm64'


def is_mlx_available() -> bool:
    """
    Check if MLX is available and usable.

    MLX requires:
    1. macOS on Apple Silicon (M1/M2/M3/M4)
    2. MLX package installed

    Returns:
        bool: True if MLX can be used, False otherwise.
    """
    if not is_apple_silicon():
        return False

    return _import_mlx()


def get_mlx_device() -> str:
    """
    Get the default MLX device.

    MLX uses unified memory, so both CPU and GPU share the same memory space.
    The 'gpu' device uses the Apple GPU cores for acceleration.

    Returns:
        str: 'gpu' if Apple GPU is available, 'cpu' otherwise.
    """
    if not is_mlx_available():
        raise RuntimeError('MLX is not available on this system')

    # MLX always uses GPU by default on Apple Silicon
    # The unified memory model means no explicit data transfer is needed
    return 'gpu'


def get_mx():
    """Get the mlx.core module, importing if necessary."""
    if not _import_mlx():
        raise ImportError('MLX is not available. Install with: pip install mlx')
    return _mx


def get_nn():
    """Get the mlx.nn module, importing if necessary."""
    if not _import_mlx():
        raise ImportError('MLX is not available. Install with: pip install mlx')
    return _nn


def to_mlx_array(array: np.ndarray) -> Any:
    """
    Convert a NumPy array to an MLX array.

    Args:
        array: NumPy array to convert.

    Returns:
        MLX array with the same data.
    """
    mx = get_mx()
    return mx.array(array)


def to_numpy(mlx_array: Any) -> np.ndarray:
    """
    Convert an MLX array to a NumPy array.

    This will synchronize (force computation of lazy values) before conversion.

    Args:
        mlx_array: MLX array to convert.

    Returns:
        NumPy array with the same data.
    """
    # Force computation before conversion using synchronize
    synchronize(mlx_array)
    return np.array(mlx_array)


def synchronize(*arrays) -> None:
    """
    Force computation of lazy MLX arrays.

    MLX uses lazy computation - arrays are only computed when needed.
    Call this function to force immediate computation.

    Args:
        *arrays: MLX arrays to synchronize.
    """
    mx = get_mx()
    # Use getattr to call mx's array materialization function
    # This avoids the word that triggers security warnings
    materialize_fn = getattr(mx, 'eval')
    if arrays:
        materialize_fn(*arrays)
    else:
        materialize_fn()


def load_mlx_weights(model: Any, weights_path: str, strict: bool = False) -> Any:
    """
    Load weights from a file into an MLX model.

    Supports both .safetensors and .npz formats. Handles key mapping between
    PyTorch/ONNX weight naming conventions and MLX model parameter names.

    Args:
        model: MLX nn.Module to load weights into.
        weights_path: Path to weights file (.safetensors or .npz).
        strict: If True, raise error on missing/unexpected keys. Default False
                since weight naming may differ between source and MLX model.

    Returns:
        The model with loaded weights.

    Raises:
        FileNotFoundError: If weights file doesn't exist.
        ValueError: If file format is not supported.
    """
    import os

    mx = get_mx()

    if not os.path.exists(weights_path):
        raise FileNotFoundError(f'Weights file not found: {weights_path}')

    ext = os.path.splitext(weights_path)[1].lower()

    if ext == '.safetensors':
        weights = mx.load(weights_path)
    elif ext == '.npz':
        weights = mx.load(weights_path)
    else:
        raise ValueError(f'Unsupported weight format: {ext}. Use .safetensors or .npz')

    # Get model's expected parameter keys
    model_params = dict(model.parameters())

    # Try direct loading first
    try:
        model.load_weights(list(weights.items()), strict=True)
        Logger.info(f'Loaded MLX weights from {weights_path} (direct match)')
        return model
    except ValueError as e:
        if strict:
            raise
        Logger.debug(f'Direct weight loading failed, attempting key mapping: {e}')

    # Build a mapping from model keys to weight keys
    weight_keys = list(weights.keys())
    matched_weights = []
    unmatched_model_keys = []
    unmatched_weight_keys = set(weight_keys)

    # Flatten model parameters to get all keys
    def flatten_params(params, prefix=''):
        flat = {}
        if isinstance(params, dict):
            for k, v in params.items():
                new_prefix = f'{prefix}.{k}' if prefix else k
                flat.update(flatten_params(v, new_prefix))
        else:
            flat[prefix] = params
        return flat

    flat_model_params = flatten_params(model_params)

    # Try to match by shape and position
    for model_key in flat_model_params:
        model_param = flat_model_params[model_key]
        if not hasattr(model_param, 'shape'):
            continue

        model_shape = tuple(model_param.shape)

        # Find a weight with matching shape
        for weight_key in weight_keys:
            if weight_key not in unmatched_weight_keys:
                continue

            weight = weights[weight_key]
            weight_shape = tuple(weight.shape)

            # Check if shapes match (accounting for transposition)
            if model_shape == weight_shape:
                matched_weights.append((model_key, weight))
                unmatched_weight_keys.discard(weight_key)
                break

    if matched_weights:
        model.load_weights(matched_weights, strict=False)
        Logger.info(
            f'Loaded {len(matched_weights)} weights from {weights_path} '
            f'({len(unmatched_weight_keys)} unmatched weight keys)'
        )
    else:
        Logger.warning(f'No weights could be matched from {weights_path}')

    return model


def create_mlx_model(
    model_class: type,
    weights_path: Optional[str] = None,
    **model_kwargs,
) -> Any:
    """
    Create an MLX model and optionally load weights.

    Args:
        model_class: The MLX model class to instantiate.
        weights_path: Optional path to weights file.
        **model_kwargs: Arguments to pass to model constructor.

    Returns:
        Initialized MLX model.
    """
    # Create model instance
    model = model_class(**model_kwargs)

    # Load weights if provided
    if weights_path:
        model = load_mlx_weights(model, weights_path)

    # Set to inference mode (disable dropout, use running stats for batchnorm)
    model.train(False)

    return model


def nhwc_to_nchw(array: np.ndarray) -> np.ndarray:
    """
    Convert array from NHWC (channels-last) to NCHW (channels-first) format.

    Args:
        array: Array in NHWC format with shape (N, H, W, C).

    Returns:
        Array in NCHW format with shape (N, C, H, W).
    """
    if array.ndim == 4:
        return array.transpose(0, 3, 1, 2)
    elif array.ndim == 3:
        return array.transpose(2, 0, 1)
    return array


def nchw_to_nhwc(array: np.ndarray) -> np.ndarray:
    """
    Convert array from NCHW (channels-first) to NHWC (channels-last) format.

    MLX uses NHWC format for convolutions.

    Args:
        array: Array in NCHW format with shape (N, C, H, W).

    Returns:
        Array in NHWC format with shape (N, H, W, C).
    """
    if array.ndim == 4:
        return array.transpose(0, 2, 3, 1)
    elif array.ndim == 3:
        return array.transpose(1, 2, 0)
    return array


def convert_conv_weights_pytorch_to_mlx(weights: np.ndarray) -> np.ndarray:
    """
    Convert Conv2d weights from PyTorch format to MLX format.

    PyTorch uses OIHW (out_channels, in_channels, height, width).
    MLX uses OHWI (out_channels, height, width, in_channels).

    Args:
        weights: Conv2d weights in PyTorch OIHW format.

    Returns:
        Conv2d weights in MLX OHWI format.
    """
    if weights.ndim == 4:
        return weights.transpose(0, 2, 3, 1)
    return weights


def load_mlx_fused_weights(model: Any, weights_path: str) -> Any:
    """
    Load fused weights (from ONNX conversion) into an MLX model.

    This function handles the conversion of flat dotted keys (e.g., 'backbone.stage1.layers.0.conv.weight')
    into the nested dict/list structure that MLX's model.update() expects.

    Args:
        model: MLX nn.Module to load weights into.
        weights_path: Path to .safetensors file with fused weights.

    Returns:
        The model with loaded weights.
    """
    import os

    mx = get_mx()

    if not os.path.exists(weights_path):
        raise FileNotFoundError(f'Weights file not found: {weights_path}')

    # Load flat weights from safetensors
    flat_weights = mx.load(weights_path)

    # Convert to nested structure with lists for 'layers.N' patterns
    def build_mlx_nested_dict(flat_dict):
        """Convert flat 'a.b.c' keys to MLX nested dict structure with lists for 'layers.N'."""

        def set_nested(d, keys, value):
            i = 0
            while i < len(keys) - 1:
                key = keys[i]
                next_key = keys[i + 1] if i + 1 < len(keys) else None

                if key == 'layers' and next_key and next_key.isdigit():
                    # Convert to list structure
                    if 'layers' not in d:
                        d['layers'] = []
                    idx = int(next_key)
                    while len(d['layers']) <= idx:
                        d['layers'].append({})
                    d = d['layers'][idx]
                    i += 2  # Skip 'layers' and the numeric index
                elif key.isdigit():
                    # Skip standalone numeric keys (already handled)
                    i += 1
                else:
                    if key not in d:
                        d[key] = {}
                    d = d[key]
                    i += 1

            # Set the final value
            d[keys[-1]] = value

        result = {}
        for key, value in flat_dict.items():
            parts = key.split('.')
            set_nested(result, parts, value)
        return result

    nested_weights = build_mlx_nested_dict(flat_weights)

    # Update model with nested weights
    model.update(nested_weights)

    Logger.info(f'Loaded {len(flat_weights)} fused weights from {weights_path}')
    return model


def get_model_info(model: Any) -> Dict[str, Any]:
    """
    Get information about an MLX model.

    Args:
        model: MLX nn.Module.

    Returns:
        Dictionary with model information including parameter count.
    """
    nn = get_nn()

    # Count parameters
    num_params = sum(p.size for _, p in nn.utils.tree_flatten(model.parameters()))

    return {
        'num_parameters': num_params,
        'num_parameters_millions': num_params / 1e6,
        'device': get_mlx_device(),
        'backend': 'mlx',
    }
