#!/usr/bin/env python3
"""
Convert MLX Gemma-4 26B to MarkBase-12B compatible format

Usage:
    python convert_mlx_26b.py --input ~/.cache/huggingface/hub/models--mlx-community--gemma-4-26b-a4b-mxfp4 --output ~/models/gemma-4-26b-standard
"""

import argparse
import json
import os
import shutil
from pathlib import Path

try:
    from safetensors.torch import load_file, save_file
    import torch
except ImportError:
    print("需要安装依赖:")
    print("  pip install safetensors torch")
    exit(1)

def convert_mlx_to_standard(input_dir: str, output_dir: str):
    """转换 MLX 26B 为标准格式"""
    
    print("=== MLX 26B → 标准 4-bit 转换 ===")
    
    input_path = Path(input_dir)
    output_path = Path(output_dir)
    
    # 创建输出目录
    output_path.mkdir(parents=True, exist_ok=True)
    
    print(f"\n步骤 1: 加载 MLX 权重")
    
    # 加载所有 shards
    weights = {}
    shard_files = [
        "model-00001-of-00003.safetensors",
        "model-00002-of-00003.safetensors",
        "model-00003-of-00003.safetensors"
    ]
    
    for shard in shard_files:
        shard_path = input_path / shard
        if shard_path.exists():
            print(f"  加载 {shard}...")
            shard_weights = load_file(str(shard_path))
            weights.update(shard_weights)
    
    print(f"  ✓ 总权重数: {len(weights)}")
    
    print(f"\n步骤 2: 重命名权重")
    
    # 重命名: language_model.model.layers.X → layers.X
    renamed_weights = {}
    
    for key, tensor in weights.items():
        # 移除 language_model.model 前缀
        new_key = key
        if key.startswith("language_model.model."):
            new_key = key.replace("language_model.model.", "")
        
        # 处理 embed_tokens (特殊情况)
        if new_key.startswith("embed_tokens"):
            # 保留原名
            pass
        
        # 处理 vision tower
        if new_key.startswith("embed_vision"):
            # 重命名为 vision_tower
            new_key = new_key.replace("embed_vision", "vision_tower")
        
        renamed_weights[new_key] = tensor
        
        if len(renamed_weights) % 100 == 0:
            print(f"  已处理 {len(renamed_weights)}/{len(weights)} 权重")
    
    print(f"  ✓ 重命名完成")
    
    print(f"\n步骤 3: 转换 scales 格式")
    
    # uint8 scales → BF16
    converted_weights = {}
    
    for key, tensor in renamed_weights.items():
        if ".scales" in key and tensor.dtype == torch.uint8:
            # uint8 → float32 → bfloat16
            converted = tensor.float().bfloat16()
            converted_weights[key] = converted
            print(f"  转换 {key}: uint8 → BF16")
        else:
            converted_weights[key] = tensor
    
    print(f"  ✓ scales 转换完成")
    
    print(f"\n步骤 4: 保存为单个 safetensors")
    
    # 合并为单个文件
    output_weights_file = output_path / "model.safetensors"
    save_file(converted_weights, str(output_weights_file))
    
    print(f"  ✓ 保存到: {output_weights_file}")
    
    print(f"\n步骤 5: 创建 config.json")
    
    # 加载原始 config
    config_path = input_path / "config.json"
    if config_path.exists():
        with open(config_path) as f:
            mlx_config = json.load(f)
        
        # 获取 text_config
        text_config = mlx_config.get("text_config", {})
        
        # 创建标准 config
        standard_config = {
            "model_type": "gemma4",
            "architectures": ["Gemma4ForConditionalGeneration"],
            "hidden_size": text_config.get("hidden_size", 2816),
            "num_hidden_layers": 42,  # 需要从权重推断
            "vocab_size": 262144,
            "intermediate_size": text_config.get("intermediate_size", 2112),
            "head_dim": text_config.get("head_dim", 256),
            "num_attention_heads": 8,
            "num_key_value_heads": 2,
            "max_position_embeddings": 131072,
            "quantization_config": {
                "bits": 4,
                "group_size": 64,
                "quant_method": "custom"
            }
        }
        
        # 添加 vision config (如果有)
        vision_config = mlx_config.get("vision_config", {})
        if vision_config:
            standard_config["vision_config"] = {
                "hidden_size": vision_config.get("hidden_size", 768),
                "num_hidden_layers": 16,
                "patch_size": 16,
                "image_size": 224
            }
        
        # 写入 config
        with open(output_path / "config.json", "w") as f:
            json.dump(standard_config, f, indent=2)
        
        print(f"  ✓ config.json 创建完成")
    
    print(f"\n步骤 6: 复制 tokenizer 文件")
    
    # 复制 tokenizer 相关文件
    tokenizer_files = [
        "tokenizer.json",
        "tokenizer_config.json",
        "generation_config.json",
        "chat_template.jinja"
    ]
    
    for file in tokenizer_files:
        src = input_path / file
        if src.exists():
            dst = output_path / file
            shutil.copy2(src, dst)
            print(f"  ✓ 复制 {file}")
    
    print(f"\n=== 转换完成 ===")
    print(f"输出目录: {output_path}")
    print(f"总权重: {len(converted_weights)}")
    
    # 显示文件列表
    print(f"\n输出文件:")
    for file in sorted(output_path.iterdir()):
        if file.is_file():
            size = file.stat().st_size / 1024 / 1024
            print(f"  {file.name}: {size:.1f} MB")
    
    print(f"\n下一步:")
    print(f"  swift run G12BServer {output_dir} 8080 gemma-26b")
    print(f"  或测试:")
    print(f"  swift test --filter test26BModelLoading")

def main():
    parser = argparse.ArgumentParser(description="转换 MLX 26B 为标准格式")
    parser.add_argument("--input", required=True, help="MLX 模型目录")
    parser.add_argument("--output", required=True, help="输出目录")
    
    args = parser.parse_args()
    
    convert_mlx_to_standard(args.input, args.output)

if __name__ == "__main__":
    main()