"""Convert ACE-Step PyTorch weights to MLX format.

Downloads from HuggingFace, remaps keys, transposes Conv1d weights,
fuses weight_norm, and saves as safetensors.

Usage:
    python -m mlx_music.models.acestep.convert --output ./weights

HuggingFace source:
    DiT: ACE-Step/Ace-Step1.5  (subdir: acestep-v15-turbo/)
         - model.safetensors  (~4.79 GB)
         - config.json
         - silence_latent.pt

    VAE: ACE-Step/Ace-Step1.5  (subdir: vae/)
         - diffusion_pytorch_model.safetensors  (~337 MB)
         - config.json

Key mapping reference:
    DiT PyTorch (AceStepConditionGenerationModel via transformers PreTrainedModel):
        decoder.proj_in.1.weight          -> decoder.proj_in.weight     (Conv1d in Sequential[1])
        decoder.proj_out.1.weight         -> decoder.proj_out.weight    (ConvTranspose1d in Sequential[1])
        decoder.layers.N.self_attn.*      -> same                       (q/k/v/o_proj, q_norm, k_norm)
        decoder.layers.N.cross_attn.*     -> same
        decoder.layers.N.mlp.*            -> same                       (gate_proj, up_proj, down_proj)
        decoder.layers.N.scale_shift_table -> same                      (nn.Parameter)
        decoder.time_embed.*              -> same                       (linear_1, linear_2, time_proj)
        decoder.time_embed_r.*            -> same
        decoder.condition_embedder.*      -> same                       (nn.Linear)
        decoder.norm_out.*                -> same
        decoder.scale_shift_table         -> same                       (nn.Parameter)
        decoder.rotary_emb.*              -> SKIP (MLX uses nn.RoPE)
        encoder.text_projector.*          -> same
        encoder.lyric_encoder.embed_tokens.* -> same
        encoder.lyric_encoder.norm.*      -> same
        encoder.lyric_encoder.layers.N.*  -> same
        encoder.lyric_encoder.rotary_emb.* -> SKIP
        encoder.timbre_encoder.*          -> same
        encoder.timbre_encoder.rotary_emb.* -> SKIP
        tokenizer.*                       -> SKIP (not needed for inference)
        detokenizer.*                     -> SKIP (not needed for inference)
        null_condition_emb                -> same                       (nn.Parameter)

    VAE PyTorch (AutoencoderOobleck via diffusers ModelMixin):
        All Conv1d/ConvTranspose1d use weight_norm, creating weight_g + weight_v pairs.
        We fuse them: weight = weight_g * weight_v / norm(weight_v)
        Then transpose: [O, I, K] -> [O, K, I]
        Snake1d alpha/beta: PyTorch [1, C, 1] -> MLX [C] (squeeze)
"""

import argparse
import json
import shutil
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from mlx_music.core.convert_utils import (
    is_conv1d_weight,
    load_safetensors,
    save_mlx_weights,
    transpose_conv1d_weight,
)

import mlx.core as mx

# HuggingFace model IDs
DIT_MODEL_ID = "ACE-Step/Ace-Step1.5"
DIT_SUBDIR = "acestep-v15-turbo"
VAE_SUBDIR = "vae"

# ---------------------------------------------------------------------------
# DiT key prefixes to skip entirely (not needed for MLX inference)
# ---------------------------------------------------------------------------
DIT_SKIP_PREFIXES = [
    # Lambda wrappers in nn.Sequential for proj_in/proj_out (indices 0 and 2)
    "decoder.proj_in.0.",
    "decoder.proj_in.2.",
    "decoder.proj_out.0.",
    "decoder.proj_out.2.",
    # RoPE embeddings — MLX uses nn.RoPE which computes these on-the-fly
    "decoder.rotary_emb.",
    "encoder.lyric_encoder.rotary_emb.",
    "encoder.timbre_encoder.rotary_emb.",
    # Tokenizer and detokenizer — not needed for inference
    "tokenizer.",
    "detokenizer.",
]

# DiT keys that are Conv1d weights needing [O,I,K] -> [O,K,I] transpose.
DIT_CONV1D_KEY_PATTERNS = [
    "decoder.proj_in.1.weight",
]

# DiT keys that are ConvTranspose1d weights needing [I,O,K] -> [O,K,I] transpose.
# PyTorch ConvTranspose1d stores weights as [in_channels, out_channels, kernel_size].
DIT_CONV_TRANSPOSE_KEY_PATTERNS = [
    "decoder.proj_out.1.weight",
]


def should_skip_dit_key(key: str) -> bool:
    """Check if a DiT weight key should be skipped entirely.

    Skips:
      - Lambda wrappers (proj_in.0, proj_in.2, proj_out.0, proj_out.2)
      - Rotary embedding tables (we use nn.RoPE)
      - Tokenizer/detokenizer (not needed for inference)
    """
    for prefix in DIT_SKIP_PREFIXES:
        if key.startswith(prefix):
            return True
    return False


def remap_dit_key(key: str) -> str:
    """Remap a PyTorch DiT state_dict key to the MLX module path.

    Remappings:
      - decoder.proj_in.1.{weight,bias}  -> decoder.proj_in.{weight,bias}
        (PyTorch wraps Conv1d in Sequential at index 1; MLX uses bare Conv1d)
      - decoder.proj_out.1.{weight,bias} -> decoder.proj_out.{weight,bias}
        (Same pattern for ConvTranspose1d)
      - All other keys pass through unchanged (the MLX module tree mirrors
        the PyTorch tree for attention, MLP, norms, embeddings, etc.)
    """
    key = key.replace("decoder.proj_in.1.", "decoder.proj_in.")
    key = key.replace("decoder.proj_out.1.", "decoder.proj_out.")
    return key


def is_dit_conv1d_weight(key: str) -> bool:
    """Check if a DiT key is a Conv1d weight needing [O,I,K] -> [O,K,I] transpose."""
    for pattern in DIT_CONV1D_KEY_PATTERNS:
        if key == pattern:
            return True
    return False


def is_dit_conv_transpose_weight(key: str) -> bool:
    """Check if a DiT key is a ConvTranspose1d weight needing [I,O,K] -> [O,K,I] transpose."""
    for pattern in DIT_CONV_TRANSPOSE_KEY_PATTERNS:
        if key == pattern:
            return True
    return False


def convert_dit_weights(weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
    """Convert DiT weights from PyTorch to MLX format.

    Processing steps for each key:
      1. Skip keys matching DIT_SKIP_PREFIXES (Lambda, rotary_emb, tokenizer, etc.)
      2. Remap Sequential indices (proj_in.1 -> proj_in, proj_out.1 -> proj_out)
      3. Transpose 3D Conv1d weights from [O, I, K] to [O, K, I]

    Args:
        weights: Dict of weight name -> mx.array (loaded via mx.load).

    Returns:
        Dict of MLX weight name -> mx.array, ready to load into the model.
    """
    mlx_weights = {}
    skipped = []
    transposed = []
    remapped = []

    for key, value in weights.items():
        # Step 1: Skip unwanted keys
        if should_skip_dit_key(key):
            skipped.append(key)
            continue

        # Step 2: Remap key path
        new_key = remap_dit_key(key)
        if new_key != key:
            remapped.append(f"{key} -> {new_key}")

        # Step 3: Transpose Conv1d weights [O, I, K] -> [O, K, I]
        if is_dit_conv1d_weight(key) and value.ndim == 3:
            value = transpose_conv1d_weight(value)
            transposed.append(f"{key} -> {new_key} {value.shape}")
        # Step 3b: Transpose ConvTranspose1d weights [I, O, K] -> [O, K, I]
        elif is_dit_conv_transpose_weight(key) and value.ndim == 3:
            if isinstance(value, mx.array):
                value = mx.transpose(value, axes=(1, 2, 0))
            else:
                value = np.transpose(value, (1, 2, 0))
            transposed.append(f"{key} -> {new_key} {value.shape} (conv_transpose)")

        mlx_weights[new_key] = value

    # Report
    if skipped:
        print(
            f"  Skipped {len(skipped)} keys (Lambda/rotary_emb/tokenizer/detokenizer)"
        )
    if remapped:
        print(f"  Remapped {len(remapped)} keys:")
        for r in remapped:
            print(f"    {r}")
    if transposed:
        print(f"  Transposed {len(transposed)} Conv1d weights:")
        for t in transposed:
            print(f"    {t}")

    return mlx_weights


# ---------------------------------------------------------------------------
# VAE weight_norm fusion
# ---------------------------------------------------------------------------


def fuse_weight_norm(
    weight_g: Union[np.ndarray, mx.array],
    weight_v: Union[np.ndarray, mx.array],
) -> Union[np.ndarray, mx.array]:
    """Fuse weight_norm parameters into a single weight tensor.

    PyTorch weight_norm decomposes a weight W into:
        W = g * (v / ||v||)
    where:
        weight_g: magnitude, shape [O, 1, 1] for Conv1d (or [O] scalar per output channel)
        weight_v: direction, shape [O, I, K] for Conv1d

    The norm is computed over all dims except dim=0 (output channels):
        ||v||_i = sqrt(sum(v[i, :, :]^2))

    Args:
        weight_g: Magnitude tensor from weight_norm, shape [O, 1, 1] or [O].
        weight_v: Direction tensor from weight_norm, shape [O, I, K].

    Returns:
        Fused weight tensor of same shape as weight_v.
    """
    use_mlx = isinstance(weight_v, mx.array)

    if use_mlx:
        # Cast to float32 for numerical stability during fusion
        weight_g = weight_g.astype(mx.float32)
        weight_v = weight_v.astype(mx.float32)

        axes_to_reduce = list(range(1, weight_v.ndim))
        v_norm = mx.sqrt(mx.sum(weight_v**2, axis=axes_to_reduce, keepdims=True))
        v_norm = mx.maximum(v_norm, 1e-12)

        if weight_g.ndim == 1:
            shape = [weight_g.shape[0]] + [1] * (weight_v.ndim - 1)
            weight_g = weight_g.reshape(shape)

        return weight_g * (weight_v / v_norm)
    else:
        axes_to_reduce = tuple(range(1, weight_v.ndim))
        v_norm = np.sqrt(np.sum(weight_v**2, axis=axes_to_reduce, keepdims=True))
        v_norm = np.maximum(v_norm, 1e-12)

        if weight_g.ndim == 1:
            shape = [weight_g.shape[0]] + [1] * (weight_v.ndim - 1)
            weight_g = weight_g.reshape(shape)

        return weight_g * (weight_v / v_norm)


def reshape_snake_param(
    value: Union[np.ndarray, mx.array],
) -> Union[np.ndarray, mx.array]:
    """Reshape Snake1d alpha/beta from PyTorch to MLX format.

    PyTorch Snake1d stores alpha and beta as [1, C, 1] (for [B, C, T] broadcasting).
    MLX Snake1d stores them as [C] (for [B, T, C] channels-last broadcasting).

    Args:
        value: Parameter tensor, expected shape [1, C, 1].

    Returns:
        Squeezed tensor of shape [C].
    """
    if isinstance(value, mx.array):
        return mx.squeeze(value)
    return value.squeeze()


def convert_vae_weights(weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
    """Convert VAE (AutoencoderOobleck) weights from PyTorch to MLX format.

    Processing steps:
      1. Collect weight_g / weight_v pairs and fuse them into single weight tensors
      2. Transpose all Conv1d/ConvTranspose1d weights [O, I, K] -> [O, K, I]
      3. Reshape Snake1d alpha/beta from [1, C, 1] to [C]

    The diffusers AutoencoderOobleck applies weight_norm to ALL Conv1d and
    ConvTranspose1d layers, so the state_dict contains:
        encoder.conv1.weight_g    [O, 1, 1]
        encoder.conv1.weight_v    [O, I, K]
        encoder.conv1.bias        [O]
    instead of a single weight tensor.

    After fusion and transpose, our MLX model expects:
        encoder.conv1.weight      [O, K, I]   (channels-last conv)
        encoder.conv1.bias        [O]

    Args:
        weights: Dict of weight name -> mx.array (loaded via mx.load).

    Returns:
        Dict of MLX weight name -> mx.array, ready to load into the model.
    """
    # Pass 1: Collect weight_g and weight_v pairs, group by base key
    weight_g_keys = {}  # base_key -> weight_g array
    weight_v_keys = {}  # base_key -> weight_v array
    other_keys = {}  # all non weight_norm keys

    for key, value in weights.items():
        if key.endswith(".weight_g"):
            base = key[: -len(".weight_g")]
            weight_g_keys[base] = value
        elif key.endswith(".weight_v"):
            base = key[: -len(".weight_v")]
            weight_v_keys[base] = value
        else:
            other_keys[key] = value

    # Pass 2: Fuse weight_norm pairs and build output
    mlx_weights = {}
    fused = []
    transposed = []
    reshaped_snake = []

    # Fuse weight_g + weight_v -> weight, then transpose
    for base_key in sorted(set(weight_g_keys.keys()) | set(weight_v_keys.keys())):
        if base_key not in weight_g_keys:
            raise ValueError(f"Found weight_v but no weight_g for: {base_key}")
        if base_key not in weight_v_keys:
            raise ValueError(f"Found weight_g but no weight_v for: {base_key}")

        g = weight_g_keys[base_key]
        v = weight_v_keys[base_key]
        fused_weight = fuse_weight_norm(g, v)
        fused.append(f"{base_key} (g={g.shape}, v={v.shape} -> {fused_weight.shape})")

        # Transpose conv weights
        if fused_weight.ndim == 3:
            if "conv_t" in base_key:
                # ConvTranspose1d: PyTorch [I, O, K] -> MLX [O, K, I]
                if isinstance(fused_weight, mx.array):
                    fused_weight = mx.transpose(fused_weight, axes=(1, 2, 0))
                else:
                    fused_weight = np.transpose(fused_weight, (1, 2, 0))
                transposed.append(
                    f"{base_key}.weight {fused_weight.shape} (conv_transpose)"
                )
            else:
                # Conv1d: PyTorch [O, I, K] -> MLX [O, K, I]
                fused_weight = transpose_conv1d_weight(fused_weight)
                transposed.append(f"{base_key}.weight {fused_weight.shape}")

        out_key = f"{base_key}.weight"
        mlx_weights[out_key] = fused_weight

    # Process remaining (non weight_norm) keys
    for key, value in other_keys.items():
        # Snake1d alpha/beta: [1, C, 1] -> [C]
        if ".alpha" in key or ".beta" in key:
            if value.ndim == 3:
                value = reshape_snake_param(value)
                reshaped_snake.append(f"{key} -> {value.shape}")

        # Any remaining 3D .weight tensors that weren't weight-normed
        elif value.ndim == 3 and key.endswith(".weight"):
            value = transpose_conv1d_weight(value)
            transposed.append(f"{key} {value.shape}")

        mlx_weights[key] = value

    # Report
    print(f"  Fused {len(fused)} weight_norm pairs:")
    for f in fused:
        print(f"    {f}")
    print(f"  Transposed {len(transposed)} Conv1d weights")
    if reshaped_snake:
        print(f"  Reshaped {len(reshaped_snake)} Snake1d params:")
        for s in reshaped_snake:
            print(f"    {s}")

    return mlx_weights


# ---------------------------------------------------------------------------
# Download and convert pipeline
# ---------------------------------------------------------------------------


def download_model(model_id: str, subdir: Optional[str] = None) -> Path:
    """Download model from HuggingFace Hub.

    Args:
        model_id: HuggingFace model ID (e.g., "ACE-Step/Ace-Step1.5")
        subdir: Subdirectory within the repo to download

    Returns:
        Path to downloaded model directory.
    """
    from huggingface_hub import snapshot_download

    allow_patterns = None
    if subdir:
        allow_patterns = [f"{subdir}/*"]

    cache_dir = snapshot_download(
        model_id,
        allow_patterns=allow_patterns,
    )
    if subdir:
        return Path(cache_dir) / subdir
    return Path(cache_dir)


def convert_all(output_dir: str, skip_download: bool = False) -> None:
    """Convert all ACE-Step components to MLX format.

    Downloads DiT and VAE weights from HuggingFace, converts them,
    and saves to output_dir with the structure:
        output_dir/
            dit/model.safetensors
            dit/config.json
            vae/model.safetensors
            vae/config.json

    Args:
        output_dir: Directory to save converted weights
        skip_download: If True, assume weights are already in HF cache
    """
    output_dir = Path(output_dir)

    # --- DiT ---
    print("Converting DiT weights...")
    dit_dir = download_model(DIT_MODEL_ID, DIT_SUBDIR)
    dit_weights = load_safetensors(dit_dir)
    print(f"  Loaded {len(dit_weights)} weight tensors from {dit_dir}")

    mlx_dit = convert_dit_weights(dit_weights)
    dit_out = output_dir / "dit"
    dit_out.mkdir(parents=True, exist_ok=True)
    save_mlx_weights(mlx_dit, str(dit_out / "model.safetensors"))
    # Copy config
    shutil.copy2(dit_dir / "config.json", dit_out / "config.json")
    print(f"  Saved {len(mlx_dit)} weights to {dit_out}")

    # --- VAE ---
    print("Converting VAE weights...")
    vae_dir = download_model(DIT_MODEL_ID, VAE_SUBDIR)
    vae_weights = load_safetensors(vae_dir)
    print(f"  Loaded {len(vae_weights)} weight tensors from {vae_dir}")

    mlx_vae = convert_vae_weights(vae_weights)
    vae_out = output_dir / "vae"
    vae_out.mkdir(parents=True, exist_ok=True)
    save_mlx_weights(mlx_vae, str(vae_out / "model.safetensors"))
    shutil.copy2(vae_dir / "config.json", vae_out / "config.json")
    print(f"  Saved {len(mlx_vae)} weights to {vae_out}")

    # --- Silence latent ---
    silence_path = dit_dir / "silence_latent.pt"
    if silence_path.exists():
        print("Converting silence_latent.pt...")
        import torch

        silence = torch.load(str(silence_path), weights_only=True).numpy()
        # PyTorch stores as [1, C, T], we want [1, T, C] (channels-last)
        if silence.ndim == 3 and silence.shape[1] < silence.shape[2]:
            silence = np.transpose(silence, (0, 2, 1))
        mx_silence = {"silence_latent": mx.array(silence)}
        save_mlx_weights(mx_silence, str(dit_out / "silence_latent.safetensors"))
        print(f"  Saved silence_latent to {dit_out}")

    print(f"\nConversion complete! Weights saved to {output_dir}")


def main():
    parser = argparse.ArgumentParser(description="Convert ACE-Step weights to MLX")
    parser.add_argument(
        "--output",
        type=str,
        default="./weights",
        help="Output directory for converted weights",
    )
    args = parser.parse_args()
    convert_all(args.output)


if __name__ == "__main__":
    main()
