import os
from pathlib import Path
import json
import cv2

from app.utils.card_inference.fry_bisenetv2_predictor_V04_250819 import FryBisenetV2Predictor

BASE_PATH = Path(__file__).parent.parent.absolute()


def predict_single_image(config_params: dict,
                         img_path: str,
                         output_dir: str,
                         only_json=False):
    # 配置参数
    model_path = BASE_PATH / config_params["pth_path"]
    real_seg_class_dict = config_params['class_dict']
    imgSize_train_dict = config_params['img_size']
    confidence = config_params['confidence']
    input_channels = config_params['input_channels']

    # 为不同类别设置不同颜色（可选）
    label_colors_dict = {
        'outer_box': (255, 0, 0),
    }

    # 创建预测器
    predictor = FryBisenetV2Predictor(
        pth_path=str(model_path),
        real_seg_class_dict=real_seg_class_dict,
        imgSize_train_dict=imgSize_train_dict,
        confidence=confidence,
        label_colors_dict=label_colors_dict,
        input_channels=input_channels,
    )

    # 单张图片预测
    print("=== 单张图片预测 ===")
    now_img_path = img_path
    answer_json_dir_str = output_dir

    if not only_json:
        result = predictor.predict_single_image(
            img_path=now_img_path,
            save_visualization=True,
            save_json=True,
            answer_json_dir_str=answer_json_dir_str
        )
    else:
        img_bgr = cv2.imread(now_img_path)
        result = predictor.predict_from_image(img_bgr)

    return result


if __name__ == '__main__':
    config = {
        "pokemon_outer_box": {
            "pth_path": "Model/pokemon_outer_box.pth",
            "class_dict": {1: 'outer_box'},
            "img_size": {'width': 1280, 'height': 1280},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "pokemon_inner_box": {
            "pth_path": "Model/pokemon_inner_box.pth",
            "class_dict": {1: 'inner_box'},
            "img_size": {'width': 1280, 'height': 1280},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "pokemon_back_corner_defect": {
            "pth_path": "Model/pokemon_back_corner_defect.pth",
            "class_dict": {
                1: 'wear', 2: 'wear_and_impact', 3: 'impact',
                4: 'damaged', 5: 'wear_and_stain',
            },
            "img_size": {'width': 512, 'height': 512},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "pokemon_front_corner_reflect_defect": {
            "pth_path": "Model/pokemon_front_corner_reflect_defect.pth",
            "class_dict": {"1": "impact", "2": "wear_and_impact", "3": "wear"},
            "img_size": {'width': 512, 'height': 512},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "pokemon_front_corner_no_reflect_defect": {
            "pth_path": "Model/pokemon_front_corner_no_reflect_defect.pth",
            "class_dict": {"1": "wear", "2": "wear_and_impact", "3": "impact", "4": "damaged"},
            "img_size": {'width': 512, 'height': 512},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "pokemon_front_face_no_reflect_defect": {
            "pth_path": "Model/pokemon_front_face_no_reflect_defect.pth",
            "class_dict": {"1": "scratch", "2": "pit", "3": "stain"},
            "img_size": {'width': 512, 'height': 512},
            "confidence": 0.5,
            "input_channels": 3,
        },
    }
    predict_single_image(config['pokemon_inner_box'],
                         img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\250805_pokemon_0001.jpg",
                         output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\inner")

    predict_single_image(config['pokemon_outer_box'],
                         img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\250805_pokemon_0001.jpg",
                         output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\outer")

    # predict_single_image(config['pokemon_back_corner_defect'],
    #                      img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\img_2.png",
    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\back_corner")

    predict_single_image(config['pokemon_front_face_no_reflect_defect'],
                         img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\250805_pokemon_0001.jpg",
                         output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\face_no_reflect")

    # result = predict_single_image(config['pokemon_front_corner_no_reflect_defect'],
    #                               img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\00006_250805_pokemon_0001_bottom_grid_r0_c5.jpg",
    #                               output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\corner_no_reflect")

    # result = predict_single_image(config['pokemon_front_face_no_reflect_defect'],
    #                      img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\250805_pokemon_0001_grid_r3_c4.jpg",
    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\face_no_reflect",
    #                      only_json=True)

    # print(result)