"""
fire_detection_inference.py

此模組提供火災檢測推論介面函式：
    inference(frame, params={})
    
當作為主程式執行時，也可以使用命令列參數測試推論。
"""

import os
import sys
import time
import argparse
import cv2
import numpy as np
import kp

# 固定路徑設定
# SCPU_FW_PATH = r'external\res\firmware\KL520\fw_scpu.bin'
# NCPU_FW_PATH = r'external\res\firmware\KL520\fw_ncpu.bin'
# MODEL_FILE_PATH = r'src\utils\models\fire_detection_520.nef'
# 若作為測試使用，預設的圖片檔案路徑（請根據實際環境調整）
# IMAGE_FILE_PATH = r'test_images\fire4.jpeg'


def preprocess_frame(frame):
    """
    將輸入的 numpy 陣列進行預處理：
      1. 調整大小至 (128, 128)
      2. 轉換為 BGR565 格式（KL520 常用格式）
    """
    if frame is None:
        raise Exception("輸入的 frame 為 None")
    
    print("預處理步驟：")
    print(f" - 原始 frame 大小: {frame.shape}")
    
    # 調整大小
    frame_resized = cv2.resize(frame, (128, 128))
    print(f" - 調整後大小: {frame_resized.shape}")
    
    # 轉換為 BGR565 格式
    # 注意：cv2.cvtColor 直接轉換到 BGR565 並非 OpenCV 標準用法，但假設此方法在 kneron SDK 下有效
    frame_bgr565 = cv2.cvtColor(frame_resized, cv2.COLOR_BGR2BGR565)
    print(" - 轉換為 BGR565 格式")
    
    return frame_bgr565


def postprocess(pre_output):
    """
    後處理函式：將模型輸出轉換為二元分類結果（這裡假設輸出為單一數值）
    """
    probability = pre_output[0]  # 假設模型輸出僅一個數值
    return probability


def inference(frame, params={}):
    """
    推論介面函式
      - frame: numpy 陣列（BGR 格式），輸入的原始影像
      - params: dict，包含額外參數，例如:
          'port_id': (int) 預設 0
          'model': (str) 模型檔案路徑，預設 MODEL_FILE_PATH
    回傳一個 dict，內容包含：
      - result: "Fire" 或 "No Fire"
      - probability: 推論信心分數
      - inference_time_ms: 推論耗時 (毫秒)
    """
    # 取得參數（若未提供則使用預設值）
    port_id = params.get('usb_port_id', 0)
    model_path = params.get('model')
    IMAGE_FILE_PATH = params.get('file_path')
    SCPU_FW_PATH = params.get('scpu_path')
    NCPU_FW_PATH = params.get('ncpu_path')

    print("Parameters received from main app:", params)
    try:
        # 1. 設備連接與初始化
        print('[連接設備]')
        device_group = kp.core.connect_devices(usb_port_ids=[port_id])
        print(' - 成功')
    
        print('[設置超時]')
        kp.core.set_timeout(device_group=device_group, milliseconds=5000)
        print(' - 成功')
    
        print('[上傳韌體]')
        kp.core.load_firmware_from_file(device_group=device_group,
                                        scpu_fw_path=SCPU_FW_PATH,
                                        ncpu_fw_path=NCPU_FW_PATH)
        print(' - 成功')
    
        print('[上傳模型]')
        model_descriptor = kp.core.load_model_from_file(device_group=device_group,
                                                        file_path=model_path)
        print(' - 成功')
    
        # 2. 圖像預處理：從 frame 轉換到符合 KL520 格式的輸入
        print('[預處理影像]')
        img_processed = preprocess_frame(frame)
    
        # 3. 建立推論描述物件
        inference_input_descriptor = kp.GenericImageInferenceDescriptor(
            model_id=model_descriptor.models[0].id,
            inference_number=0,
            input_node_image_list=[
                kp.GenericInputNodeImage(
                    image=img_processed,
                    image_format=kp.ImageFormat.KP_IMAGE_FORMAT_RGB565,
                    resize_mode=kp.ResizeMode.KP_RESIZE_ENABLE,
                    padding_mode=kp.PaddingMode.KP_PADDING_CORNER,
                    normalize_mode=kp.NormalizeMode.KP_NORMALIZE_KNERON
                )
            ]
        )
    
        # 4. 執行推論
        print('[執行推論]')
        start_time = time.time()
        kp.inference.generic_image_inference_send(
            device_group=device_group,
            generic_inference_input_descriptor=inference_input_descriptor
        )
        generic_raw_result = kp.inference.generic_image_inference_receive(
            device_group=device_group
        )
        inference_time = (time.time() - start_time) * 1000  # 毫秒
        print(f' - 推論耗時: {inference_time:.2f} ms')
    
        # 5. 處理推論結果
        print('[處理結果]')
        inf_node_output_list = []
        for node_idx in range(generic_raw_result.header.num_output_node):
            inference_float_node_output = kp.inference.generic_inference_retrieve_float_node(
                node_idx=node_idx,
                generic_raw_result=generic_raw_result,
                channels_ordering=kp.ChannelOrdering.KP_CHANNEL_ORDERING_CHW
            )
            inf_node_output_list.append(inference_float_node_output.ndarray.copy())
    
        # 整理成一維陣列並後處理
        probability = postprocess(np.array(inf_node_output_list).flatten())
        result_str = "Fire" if probability > 0.5 else "No Fire"
    
        # 6. 斷開設備連接
        kp.core.disconnect_devices(device_group=device_group)
        print('[已斷開設備連接]')
    
        # 回傳結果
        return {
            "result": result_str,
            "probability": probability,
            "inference_time_ms": inference_time
        }
    
    except Exception as e:
        print(f"錯誤: {str(e)}")
        # 嘗試斷開設備（若有連線）
        try:
            kp.core.disconnect_devices(device_group=device_group)
        except Exception:
            pass
        raise


# 若作為主程式執行，支援從命令列讀取圖片檔案並測試推論
# if __name__ == '__main__':
#     parser = argparse.ArgumentParser(
#         description='KL520 Fire Detection Model Inference'
#     )
#     parser.add_argument(
#         '-p', '--port_id', help='Port ID (Default: 0)', default=0, type=int
#     )
#     parser.add_argument(
#         '-m', '--model', help='NEF model path', default=model_path, type=str
#     )
#     parser.add_argument(
#         '-i', '--img', help='Image path', default=IMAGE_FILE_PATH, type=str
#     )
#     args = parser.parse_args()
    
#     # 讀取圖片（使用 cv2 讀取）
#     test_image = cv2.imread(args.img)
#     if test_image is None:
#         print(f"無法讀取圖片: {args.img}")
#         sys.exit(1)
    
#     # 構造參數字典
#     params = {
#         "port_id": args.port_id,
#         "model": args.model
#     }
    
#     # 呼叫推論介面函式
#     result = inference(test_image, params)
    
#     print("\n結果摘要:")
#     print(f"預測結果: {result['result']}")
#     print(f"信心分數: {result['probability']:.4f}")
#     print(f"推論時間: {result['inference_time_ms']:.2f} ms")