"""Shared weight conversion utilities for PyTorch -> MLX."""

import numpy as np
from pathlib import Path
from typing import Dict, Optional, Union

import mlx.core as mx


def transpose_conv1d_weight(
    w: Union[np.ndarray, mx.array],
) -> Union[np.ndarray, mx.array]:
    """Transpose Conv1d weight from PyTorch to MLX format.

    PyTorch Conv1d: [out_channels, in_channels, kernel_size]
    MLX Conv1d:     [out_channels, kernel_size, in_channels]
    """
    if w.ndim != 3:
        raise ValueError(f"Expected 3D tensor for Conv1d weight, got {w.ndim}D")
    if isinstance(w, mx.array):
        return mx.transpose(w, axes=(0, 2, 1))
    return np.transpose(w, (0, 2, 1))


def is_conv1d_weight(key: str, shape: tuple) -> bool:
    """Heuristic: detect Conv1d weight tensors by key name and shape.

    Conv1d weights are 3D tensors in known conv paths.
    Bias tensors are 1D and should NOT be transposed.
    """
    if len(shape) != 3:
        return False
    # Known Conv1d weight patterns in ACE-Step
    conv_patterns = [
        ".conv1.",
        ".conv2.",
        ".conv_t1.",
        "proj_in.1.",
        "proj_out.1.",  # DiT patch embed/unembed
        # VAE patterns
        ".block.",
        "conv1d",
    ]
    for pattern in conv_patterns:
        if pattern in key and key.endswith(".weight"):
            return True
    # Fallback: any 3D .weight tensor is likely a conv weight
    if key.endswith(".weight"):
        return True
    return False


def load_safetensors(path: str) -> Dict[str, mx.array]:
    """Load weights from safetensors file(s).

    Uses mx.load() which handles bfloat16 natively (NumPy doesn't support bf16).

    Args:
        path: Path to a .safetensors file or directory containing them.

    Returns:
        Dictionary of weight name -> mx.array.
    """
    path = Path(path)
    weights = {}

    if path.is_file():
        files = [path]
    elif path.is_dir():
        files = sorted(path.glob("*.safetensors"))
        if not files:
            raise FileNotFoundError(f"No .safetensors files found in {path}")
    else:
        raise FileNotFoundError(f"Path does not exist: {path}")

    for f in files:
        loaded = mx.load(str(f))
        weights.update(loaded)

    return weights


def save_mlx_weights(weights: Dict[str, mx.array], output_path: str) -> None:
    """Save MLX weights as safetensors.

    Args:
        weights: Dictionary of weight name -> mx.array.
        output_path: Path to save the .safetensors file.
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    mx.save_safetensors(str(output_path), weights)
