import os
import cv2
import numpy as np

# =========================================================
# 全局配置
# =========================================================

TEMPLATE_DIR = "./muban_image1"

# 三个 ROI（x, y, w, h）
ROI_1 = (782, 614, 164, 128)
ROI_2 = (837, 791, 100, 99)
ROI_3 = (873, 736, 141, 110)

USE_GRAY = False
TM_METHOD = cv2.TM_CCOEFF_NORMED
SCORE_THRESH = 0.4


# =========================================================
# 模板加载
# =========================================================
def load_templates(template_dir):
    templates = []
    if not os.path.isdir(template_dir):
        raise RuntimeError(f"模板目录不存在: {template_dir}")

    for fname in os.listdir(template_dir):
        if not fname.lower().endswith((".png", ".jpg", ".jpeg", ".bmp")):
            continue

        path = os.path.join(template_dir, fname)
        img = cv2.imread(path)
        if img is None:
            continue

        if USE_GRAY:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        h, w = img.shape[:2]
        name = os.path.splitext(fname)[0]

        templates.append({
            "name": name,
            "img": img,
            "h": h,
            "w": w
        })

    print(f"[INFO] 已加载 {len(templates)} 个模板")
    return templates


TEMPLATES = load_templates(TEMPLATE_DIR)


# =========================================================
# ROI 裁剪
# =========================================================
def crop_roi(img, roi):
    x, y, w, h = roi
    return img[y:y + h, x:x + w]


# =========================================================
# 单 ROI 模板匹配
# =========================================================
def match_one_roi(img, roi):
    roi_img = crop_roi(img, roi)
    if roi_img.size == 0:
        return None, 0.0

    if USE_GRAY:
        roi_img = cv2.cvtColor(roi_img, cv2.COLOR_BGR2GRAY)

    best_name = None
    best_score = -1.0

    for t in TEMPLATES:
        if t["h"] > roi_img.shape[0] or t["w"] > roi_img.shape[1]:
            continue

        res = cv2.matchTemplate(roi_img, t["img"], TM_METHOD)
        _, max_val, _, _ = cv2.minMaxLoc(res)

        if max_val > best_score:
            best_score = max_val
            best_name = t["name"]

    if best_score < SCORE_THRESH:
        return None, best_score

    return best_name, best_score


# =========================================================
# 核心逻辑：ROI_1 → ROI_2 → ROI_3
# =========================================================
def match_image(img):
    roi_list = [
        ("roi1", ROI_1),
        ("roi2", ROI_2),
        ("roi3", ROI_3),
    ]

    for roi_name, roi in roi_list:
        name, score = match_one_roi(img, roi)
        if name is not None:
            return {
                "result": name,
                "roi": roi_name,
                "score": score
            }

    return {
        "result": None,
        "roi": None,
        "score": 0.0
    }


# =========================================================
# 可视化（只画命中的 ROI）
# =========================================================
def draw_result(img, result):
    vis = img.copy()

    roi_map = {
        "roi1": ROI_1,
        "roi2": ROI_2,
        "roi3": ROI_3
    }

    if result["roi"] not in roi_map:
        return vis

    x, y, w, h = roi_map[result["roi"]]
    text = f'{result["result"]} {result["score"]:.2f}'

    cv2.rectangle(vis, (x, y), (x + w, y + h), (0, 255, 0), 2)
    cv2.putText(
        vis, text, (x, y - 10),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.7, (0, 255, 0), 2
    )

    return vis


# =========================================================
# main 测试
# =========================================================
if __name__ == "__main__":
    img = cv2.imread("1.png")
    if img is None:
        raise RuntimeError("测试图片未找到")

    result = match_image(img)
    print("[RESULT]", result)

    vis = draw_result(img, result)
    cv2.imshow("template match", vis)
    cv2.waitKey(0)
    cv2.destroyAllWindows()