import os
import numpy as np
from hailo_sdk_client import ClientRunner

from utils import make_quant_dataset


class HailoConverter:
    def __init__(self, model_name, onnx_path, hw_arch="hailo8l"):
        self.model_name = model_name
        self.onnx_path = onnx_path
        self.hw_arch = hw_arch
        self.runner = ClientRunner(hw_arch=self.hw_arch)
        self.output_dir = os.path.dirname(onnx_path)
        
        if not os.path.isfile(self.onnx_path):
            raise FileNotFoundError(f"ONNX file not found at {self.onnx_path}")

    def data_transform(self, img):
        img /= 255.0
        return img

    def prepare_dataset(self, dataset_glob, img_size=224, num_imgs=1000):
        print("\n" + "!"*50)
        print("💡 提醒: 不同種類的模型可能需要不同的預處理 (Data Transform)。")
        print("   例如: YOLO 通常僅需要 / 255.0，但有些模型可能需要減去均值或特定縮放。")
        print("   目前的預設行為是: img / 255.0")
        print("   如有需要，您可以在啟動前透過 Monkey Patch 取代 `data_transform` 方法。")
        print("   或是去nxva/nxengine/hailo/convert/calib_transform.py中找尋其他方法。")
        print("!"*50)

        confirm = input("\n確認目前的 Data Transform 邏輯正確嗎？(Enter 繼續 / n 結束程式): ").lower()
        if confirm == 'n':
            print("程式已由使用者終止。請修改預處理邏輯後再試。")
            sys.exit(0)

        print(f"\n=== 準備校準數據集: {dataset_glob} ===")
        calib_dataset, num_calib = make_quant_dataset(
            dataset_path=dataset_glob, 
            img_size=img_size, 
            output_dir=self.output_dir, 
            num_imgs=num_imgs
        )
        
        calib_dataset = calib_dataset.astype(np.float32)
        calib_dataset = self.data_transform(calib_dataset)
        
        print(f"數據集形狀: {calib_dataset.shape}")
        return calib_dataset, num_calib

    def _get_interactive_model_script(self, num_calib):
        """互動式生成模型優化脚本指令"""
        commands = [f"model_optimization_config(calibration, calibset_size={num_calib})\n"]
        
        print("\n" + "="*50)
        print("   Hailo 模型優化脚本配置 (互動環節)")
        print("="*50)
        print("必選指令: model_optimization_config (基礎校準已加入)")
        
        opt_level = 0
        comp_level = 1

        # 選項 1: Optimization Level
        print("\n--- 1. 調整優化層級 (Optimization Level) ---")
        print("用途:")
        print("  - 提高優化層級 (Level 4): 執行進階演算法 (如 AdaRound) 以恢復精度，但編譯時間較長。")
        print("  - 預設層級 (Level 0): 基本優化。")
        choice = input("是否啟用優化Level 4? (y/n, 預設 y): ").lower()
        if choice != 'n':
            opt_level = 4

        # 選項 2: Compression Level
        print("\n--- 2. 調整壓縮層級 (Compression Level) ---")
        print("用途:")
        print("  - 預設壓縮 (Level 1): 將部分權重轉為 4-bit 以提升效能。")
        print("  - 禁用壓縮 (Level 0): 若精度不足，應設為 0 以保持精度。")
        choice = input("是否禁用 4-bit 壓縮 (設定為Level 0)? (y/n, 預設 y): ").lower()
        if choice != 'n':
            comp_level = 0
        
        commands.append(f"model_optimization_flavor(compression_level={comp_level}, optimization_level={opt_level})\n")

        # 選項 3: Activation Clipping
        print("\n--- 3. 裁剪離群值 (Activation Clipping) ---")
        print("用途:")
        print("  - 裁剪活化函數輸出中的異常極端值，縮小量化間隔，顯著改善量化後的雜訊。")
        choice = input("是否啟用活化值裁剪? (y/n, 預設 y): ").lower()
        if choice != 'n':
            commands.append("pre_quantization_optimization(activation_clipping, layers={*}, mode=percentile, clipping_values=[0.01, 99.99])\n")

        # 選項 3: Fine-tune
        print("\n--- 3. 量化後微調 (Post-Quantization Fine-tuning) ---")
        print("用途:")
        print("  - 在全精度優化後進行額外的微調訓練 (Fine-tune)，適合對精度要求極高的場景。")
        choice = input("是否啟用微調 (epochs=8)? (y/n, 預設 n): ").lower()
        if choice == 'y':
            commands.append(f"post_quantization_optimization(finetune, policy=enabled, epochs=8, learning_rate=0.0001, dataset_size={num_calib})\n")

        # 選項 5: OSNet Global Avgpool Reduction
        print("\n--- 5. OSNet 特殊優化 (Global Avgpool Reduction) ---")
        print("用途:")
        print("  - 針對 OSNet 模型，若在編譯階段報錯，通常需要降低全局平均池化層的數值範圍。")
        choice = input("是否啟用 OSNet 預設池化優化 (avgpool5/7)? (y/n, 預設 n): ").lower()
        if choice == 'y':
            commands.append("pre_quantization_optimization(global_avgpool_reduction, layers=model/avgpool5, division_factors=[4, 4])\n")
            commands.append("pre_quantization_optimization(global_avgpool_reduction, layers=model/avgpool7, division_factors=[4, 4])\n")
            print("💡 提醒: 如果後續仍報錯，可能需要手動指定其他層 (layers)，請參考 nxva/nxengine/hailo/convert 中的具體實現。")

        print("\n" + "-"*50)
        print("最終生成的模型脚本:")
        for idx, cmd in enumerate(commands):
            print(f"  {idx+1}. {cmd.strip()}")
        print("-"*50 + "\n")
        
        return commands

    def convert(self, calib_dataset, num_calib, start_node_names=None, end_node_names=None, model_script_commands=None):
        if start_node_names is None:
            start_node_names = []
        
        # 步驟1: 轉換模型
        print(f"=== 步驟1: 轉換 ONNX 模型 ===")
        self.runner.translate_onnx_model(
            self.onnx_path, 
            start_node_names=start_node_names, 
            end_node_names=end_node_names
        )

        # 步驟2: 加載模型脚本 (若未提供則進入互動模式)
        if model_script_commands is None:
            model_script_commands = self._get_interactive_model_script(num_calib)
        
        self.runner.load_model_script("".join(model_script_commands))

        # 步驟3: 應用模型脚本更改
        print("=== 步驟3: 進行全精度優化 (Full Precision Optimization) ===")
        self.runner.optimize_full_precision()

        # 步驟4: 量化優化 (Quantization)
        print("=== 步驟4: 量化優化 (Quantization) ===")
        self.runner.optimize(calib_dataset)

        # 步驟5: 保存量化模型 (*.har)
        print("=== 步驟5: 保存量化模型 (*.har) ===")
        quantized_har_path = os.path.join(self.output_dir, f"{self.model_name}_quantized.har")
        self.runner.save_har(quantized_har_path)
        print(f"量化模型已保存到: {quantized_har_path}")

        # 步驟6: 編譯成HEF (*.hef)
        print("=== 步驟6: 編譯成HEF (*.hef) ===")
        hef_bytes = self.runner.compile()
        hef_path = os.path.join(self.output_dir, f"{self.model_name}.hef")
        with open(hef_path, "wb") as f:
            f.write(hef_bytes)
        print(f"HEF文件已保存到: {hef_path}")

        print("\n🎉 轉換完成！")
        return quantized_har_path, hef_path