#!/usr/bin/env python3
"""
Convert WeSpeaker PyTorch ResNet models to MLX format.

Usage:
    python convert_wespeaker_to_mlx.py [--model MODEL_ID] [--output OUTPUT_PATH]

Examples:
    # Default: convert ResNet34-LM
    python convert_wespeaker_to_mlx.py

    # Custom model
    python convert_wespeaker_to_mlx.py --model Wespeaker/wespeaker-voxceleb-resnet34-LM

    # Custom output
    python convert_wespeaker_to_mlx.py --output ./my_model_mlx.npz

Requirements:
    pip install torch numpy huggingface_hub

Notes:
    - Downloads the PyTorch avg_model from HuggingFace
    - Transposes Conv2d weights: PyTorch (O,I,H,W) → MLX (O,H,W,I)
    - Remaps key names for MLX nn.Sequential (.layers.N.) and shortcut paths
    - Skips projection.weight (not used in embedding extraction)
    - Skips num_batches_tracked (not needed for inference)
"""

import argparse
import os
import re

import numpy as np
import torch
from huggingface_hub import hf_hub_download


def convert(model_id: str, output_path: str) -> None:
    """Convert a WeSpeaker PyTorch model to MLX npz format."""
    
    print(f"Downloading PyTorch model from {model_id}...")
    pt_path = hf_hub_download(model_id, "avg_model")
    pt_state = torch.load(pt_path, map_location="cpu")
    
    print(f"Converting {len(pt_state)} parameters...")
    save = {}
    skipped = []
    
    for key, tensor in pt_state.items():
        val = tensor.numpy()
        new_key = key
        
        # 1. Map nn.Sequential indices: layer1.0. → layer1.layers.0.
        new_key = re.sub(r"(layer[1-4])\.(\d+)\.", r"\1.layers.\2.", new_key)
        
        # 2. Map shortcut paths: shortcut.0 → shortcut_conv, shortcut.1 → shortcut_bn
        new_key = new_key.replace(".shortcut.0.", ".shortcut_conv.")
        new_key = new_key.replace(".shortcut.1.", ".shortcut_bn.")
        
        # 3. Map FC layer: seg_1 → fc
        new_key = new_key.replace("seg_1.", "fc.")
        
        # 4. Skip unused keys
        if "projection" in new_key:
            skipped.append(key)
            continue
        if "num_batches_tracked" in new_key:
            skipped.append(key)
            continue
        
        # 5. Transpose Conv2d weights: PyTorch (O,I,H,W) → MLX (O,H,W,I)
        if "weight" in key and val.ndim == 4:
            val = np.transpose(val, (0, 2, 3, 1))
        
        save[new_key] = val
    
    np.savez(output_path, **save)
    size_mb = os.path.getsize(output_path) / 1024 / 1024
    
    print(f"\n✅ Saved {len(save)} params → {output_path} ({size_mb:.1f}MB)")
    if skipped:
        print(f"⏭️  Skipped {len(skipped)} keys: {skipped}")
    
    # Verification summary
    conv_keys = [k for k in save if "conv" in k and "weight" in k]
    bn_keys = [k for k in save if "bn" in k]
    fc_keys = [k for k in save if "fc" in k]
    print(f"\nBreakdown: {len(conv_keys)} conv weights, {len(bn_keys)} BN params, {len(fc_keys)} FC params")


def main():
    parser = argparse.ArgumentParser(description="Convert WeSpeaker PyTorch → MLX")
    parser.add_argument(
        "--model", 
        default="Wespeaker/wespeaker-voxceleb-resnet34-LM",
        help="HuggingFace model ID (default: Wespeaker/wespeaker-voxceleb-resnet34-LM)"
    )
    parser.add_argument(
        "--output",
        default=None,
        help="Output npz path (default: <model_name>_mlx.npz in current dir)"
    )
    args = parser.parse_args()
    
    if args.output is None:
        name = args.model.split("/")[-1].replace("wespeaker-voxceleb-", "")
        args.output = f"{name}_mlx.npz"
    
    convert(args.model, args.output)


if __name__ == "__main__":
    main()