#!/usr/bin/env python3
"""
将 YOLOv8 ONNX 模型转换为 RKNN 格式
适用于 RK3588 / RK3568 / RK3576 等平台

环境要求:
    - Ubuntu x86_64 / Docker
    - Python 3.8 / 3.9 / 3.10 / 3.11
    - rknn-toolkit2 (pip install rknn-toolkit2==2.2.0)

使用方法:
    # FP16 模式（推荐，速度快精度高）
    python 04_convert_rknn.py best.onnx -o shoe_detector.rknn -t rk3588
    
    # INT8 量化（模型更小，需要校准数据集）
    python 04_convert_rknn.py best.onnx -o shoe_detector.rknn -t rk3588 -q -d dataset.txt

支持的 target_platform:
    - rk3588 / rk3588s
    - rk3568 / rk3566
    - rk3576
    - rv1106 / rv1103 / rv1103b
    - rv1126
"""

import argparse
import os
import sys
from pathlib import Path


def check_environment():
    """检查运行环境"""
    try:
        from rknn.api import RKNN
        print("✓ RKNN Toolkit2 已安装")
        return True
    except ImportError:
        print("✗ 错误: 未安装 RKNN Toolkit2")
        print("\n请安装:")
        print("  pip install rknn-toolkit2==2.2.0")
        print("\n或从源码安装:")
        print("  https://github.com/airockchip/rknn-toolkit2")
        return False


def create_sample_dataset(onnx_path: str, output_path: str = "dataset.txt", num_samples: int = 20):
    """
    创建示例量化校准数据集
    用于 INT8 量化时提供校准图片路径
    """
    print(f"\n创建示例校准数据集: {output_path}")
    print("注意: 请用实际图片替换这些示例路径")
    
    sample_content = f"""# RKNN INT8 量化校准数据集
# 每行一个图片路径，建议使用 20-100 张典型场景图片
# 图片格式: JPG, PNG, BMP 等

# 示例路径（请替换为实际路径）:
# /path/to/train/images/img001.jpg
# /path/to/train/images/img002.jpg
# /path/to/valid/images/img001.jpg

# 提示:
# 1. 图片应与实际部署场景相似
# 2. 包含各种光照、角度、背景的样本
# 3. 建议 20-100 张，越多越慢但可能更准
"""
    
    with open(output_path, 'w') as f:
        f.write(sample_content)
    
    print(f"✓ 示例数据集已创建: {output_path}")
    print("  请编辑此文件，添加实际的图片路径")
    return output_path


def convert_onnx_to_rknn(
    onnx_path: str,
    output_path: str = None,
    target_platform: str = "rk3588",
    do_quantization: bool = False,
    dataset_path: str = None,
    verbose: bool = True
):
    """
    转换 ONNX 模型到 RKNN
    
    Args:
        onnx_path: ONNX 模型文件路径
        output_path: 输出 RKNN 文件路径，默认与 ONNX 同名
        target_platform: 目标平台，默认 rk3588
        do_quantization: 是否启用 INT8 量化
        dataset_path: 量化校准数据集路径（txt 文件，每行一张图片路径）
        verbose: 是否打印详细信息
    """
    if output_path is None:
        output_path = onnx_path.replace(".onnx", ".rknn")
    
    # 确保输出目录存在
    output_dir = os.path.dirname(output_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    print("="*70)
    print(f"ONNX 转 RKNN")
    print("="*70)
    print(f"输入:  {onnx_path}")
    print(f"输出:  {output_path}")
    print(f"目标:  {target_platform}")
    print(f"量化:  {'INT8' if do_quantization else 'FP16 (无量化)'}")
    if do_quantization:
        print(f"校准:  {dataset_path}")
    print("="*70)
    
    # 检查输入文件
    if not os.path.exists(onnx_path):
        print(f"\n✗ 错误: 找不到 ONNX 文件: {onnx_path}")
        return False
    
    # 检查数据集（如果需要量化）
    if do_quantization:
        if dataset_path is None:
            print("\n✗ 错误: INT8 量化需要提供校准数据集")
            print("  使用 --dataset 指定数据集文件路径")
            print("  或运行 --create-dataset 创建示例")
            return False
        if not os.path.exists(dataset_path):
            print(f"\n✗ 错误: 找不到数据集文件: {dataset_path}")
            return False
    
    from rknn.api import RKNN
    
    # 创建 RKNN 对象
    rknn = RKNN(verbose=verbose)
    
    try:
        # 配置模型
        print("\n[1/4] 配置模型...")
        rknn.config(
            mean_values=[[0, 0, 0]],           # YOLOv8 使用 0-255 输入
            std_values=[[255, 255, 255]],      # 归一化到 0-1
            target_platform=target_platform
        )
        print("      ✓ 完成")
        
        # 加载 ONNX
        print("\n[2/4] 加载 ONNX 模型...")
        ret = rknn.load_onnx(model=onnx_path)
        if ret != 0:
            print("      ✗ 加载失败!")
            return False
        print("      ✓ 完成")
        
        # 构建模型
        print("\n[3/4] 构建 RKNN 模型...")
        if do_quantization:
            print(f"      使用 INT8 量化，校准数据集: {dataset_path}")
            ret = rknn.build(do_quantization=True, dataset=dataset_path)
        else:
            print("      使用 FP16 模式（无量化）")
            ret = rknn.build(do_quantization=False)
        
        if ret != 0:
            print("      ✗ 构建失败!")
            return False
        print("      ✓ 完成")
        
        # 导出 RKNN
        print("\n[4/4] 导出 RKNN 模型...")
        ret = rknn.export_rknn(output_path)
        if ret != 0:
            print("      ✗ 导出失败!")
            return False
        print("      ✓ 完成")
        
    finally:
        rknn.release()
    
    # 验证输出
    if os.path.exists(output_path):
        size_mb = os.path.getsize(output_path) / (1024 * 1024)
        print("\n" + "="*70)
        print(f"✓ 转换成功!")
        print(f"  输出文件: {output_path}")
        print(f"  文件大小: {size_mb:.2f} MB")
        print("="*70)
        
        print("\n下一步:")
        print(f"  1. 复制到 RK3588:")
        print(f"     scp {output_path} orangepi@<rk3588_ip>:/home/orangepi/apps/OrangePi3588Media/models/")
        print(f"  2. 更新配置文件中的模型路径")
        return True
    else:
        print("\n✗ 错误: 输出文件未生成")
        return False


def main():
    parser = argparse.ArgumentParser(
        description="将 YOLOv8 ONNX 模型转换为 RKNN",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # FP16 模式（推荐）
  python 04_convert_rknn.py best.onnx -o shoe_detector.rknn
  
  # 指定目标平台
  python 04_convert_rknn.py best.onnx -t rk3568
  
  # INT8 量化
  python 04_convert_rknn.py best.onnx -q -d dataset.txt
  
  # 创建示例校准数据集
  python 04_convert_rknn.py --create-dataset
        """
    )
    
    parser.add_argument("onnx", nargs="?", help="ONNX 模型文件路径")
    parser.add_argument("-o", "--output", help="输出 RKNN 文件路径")
    parser.add_argument("-t", "--target", default="rk3588",
                       help="目标平台 (默认: rk3588)")
    parser.add_argument("-q", "--quantize", action="store_true",
                       help="启用 INT8 量化")
    parser.add_argument("-d", "--dataset", help="量化校准数据集路径 (txt 文件)")
    parser.add_argument("--create-dataset", action="store_true",
                       help="创建示例校准数据集并退出")
    parser.add_argument("-v", "--verbose", action="store_true", default=True,
                       help="显示详细信息")
    
    args = parser.parse_args()
    
    # 创建示例数据集
    if args.create_dataset:
        create_sample_dataset("best.onnx", "dataset.txt")
        return 0
    
    # 检查参数
    if args.onnx is None:
        parser.print_help()
        print("\n错误: 请提供 ONNX 文件路径")
        return 1
    
    # 检查环境
    if not check_environment():
        return 1
    
    # 执行转换
    success = convert_onnx_to_rknn(
        onnx_path=args.onnx,
        output_path=args.output,
        target_platform=args.target,
        do_quantization=args.quantize,
        dataset_path=args.dataset,
        verbose=args.verbose
    )
    
    return 0 if success else 1


if __name__ == "__main__":
    sys.exit(main())
