#!/usr/bin/env python3
"""
ArcFace ONNX to TFLite 转换脚本

使用方法:
    1. 先安装依赖：pip install onnx onnxruntime tensorflow
    2. 运行转换：python convert_model.py
"""

import os
import numpy as np
import subprocess
import sys

def install_requirements():
    """安装必要的包"""
    packages = ['onnx', 'onnxruntime', 'tensorflow', 'tf2onnx', 'onnx-tf']
    print("检查并安装依赖包...")
    for pkg in packages:
        try:
            __import__(pkg.replace('-', '_'))
            print(f"  ✓ {pkg} 已安装")
        except ImportError:
            print(f"  安装 {pkg}...")
            subprocess.check_call([sys.executable, '-m', 'pip', 'install', pkg, '-q'])
    print("依赖包安装完成！\n")

def download_model():
    """下载 ArcFace 模型"""
    model_url = "https://drive.google.com/uc?export=download&id=1gnt6P3jaiwfevV4hreWHPu0Mive5VRyP"
    output_path = "/Users/liushuming/projects/app/app/src/main/assets/arcface_ir50_glint360k.onnx"

    print("下载 ArcFace 模型 (IResNet-50, Glint360K 数据集)...")
    print(f"目标路径：{output_path}")

    import urllib.request
    urllib.request.urlretrieve(model_url, output_path)

    size_mb = os.path.getsize(output_path) / 1024 / 1024
    print(f"下载完成！模型大小：{size_mb:.2f} MB\n")

    return output_path

def convert_onnx_to_tflite(onnx_path, output_path):
    """将 ONNX 模型转换为 TFLite"""
    print(f"转换 ONNX -> TFLite...")
    print(f"输入：{onnx_path}")
    print(f"输出：{output_path}")

    try:
        import onnx
        from onnx_tf.backend import prepare

        # 加载 ONNX 模型
        onnx_model = onnx.load(onnx_path)
        onnx.checker.check_model(onnx_model)

        # 打印模型信息
        inputs = onnx_model.graph.input
        outputs = onnx_model.graph.output
        print(f"\n模型输入：{inputs[0].name}, 形状：{[d.dim_value for d in inputs[0].type.tensor_type.shape.dim]}")
        print(f"模型输出：{outputs[0].name}, 形状：{[d.dim_value for d in outputs[0].type.tensor_type.shape.dim]}")

        # 转换为 TensorFlow
        print("\n转换为 TensorFlow 格式...")
        tf_backend = prepare(onnx_model)
        tf_graph = tf_backend.tf_graph

        # 保存为 SavedModel
        import tensorflow as tf
        saved_model_dir = output_path.replace('.tflite', '_saved_model')

        # 清理已存在的目录
        import shutil
        if os.path.exists(saved_model_dir):
            shutil.rmtree(saved_model_dir)

        # 保存 graph
        with tf.Graph().as_default() as graph:
            tf.import_graph_def(tf_graph, name="")
            tf.saved_model.save(tf.keras.models.Model(inputs=graph.get_tensor_by_name('input:0'),
                                                      outputs=graph.get_tensor_by_name('output:0')),
                               saved_model_dir)

        # 转换为 TFLite
        print("转换为 TFLite 格式...")
        converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
        converter.optimizations = [tf.lite.Optimize.DEFAULT]

        tflite_model = converter.convert()

        with open(output_path, 'wb') as f:
            f.write(tflite_model)

        # 清理临时文件
        shutil.rmtree(saved_model_dir)

        size_mb = os.path.getsize(output_path) / 1024 / 1024
        print(f"\n✓ 转换完成！")
        print(f"  TFLite 模型：{output_path}")
        print(f"  模型大小：{size_mb:.2f} MB")

        return True

    except ImportError as e:
        print(f"\n错误：缺少依赖包 - {e}")
        print("请运行：pip install onnx-tf tensorflow")
        return False
    except Exception as e:
        print(f"\n转换错误：{e}")
        return False

def main():
    print("=" * 60)
    print("ArcFace 模型下载和转换工具")
    print("=" * 60)

    # 安装依赖
    install_requirements()

    # 下载模型
    onnx_path = download_model()

    # 转换为 TFLite
    output_path = "/Users/liushuming/projects/app/app/src/main/assets/facenet.tflite"
    convert_onnx_to_tflite(onnx_path, output_path)

    print("\n" + "=" * 60)
    print("完成！")
    print("模型已保存到：app/src/main/assets/facenet.tflite")
    print("=" * 60)

if __name__ == "__main__":
    main()