"""方法A基线：蒙版多尺度模板匹配（NCC）。

对每对样本：
- 单尺度 1.0：无蒙版 / 有蒙版 —— 验证"裸匹配是否够用"
- 多尺度 0.50~1.50（步长 0.05，有蒙版）：记录每尺度最佳分数，看目标是否为"唯一精确匹配峰"
- 可视化：最佳匹配框 + 响应热力图 → try/out/A/
"""
from pathlib import Path
import cv2
import numpy as np

ROOT = Path(__file__).resolve().parent.parent
CAP = ROOT / "captchas"
OUT = ROOT / "try" / "out" / "A"
OUT.mkdir(parents=True, exist_ok=True)

SCALES = np.arange(0.50, 1.51, 0.05)


def load_mark(path):
    """返回 (BGR, alpha)，缺 alpha 时返回全 255。"""
    try:
        m = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
        if m is None or m.ndim < 2:
            return None, None
        if m.ndim == 3 and m.shape[2] == 4:
            return m[..., :3].copy(), m[..., 3].copy()
        return m.copy(), np.full(m.shape[:2], 255, np.uint8)
    except Exception as e:
        print(f"  mark 读取失败: {e}")
        return None, None


def match_masked(big, bgr, alpha, scale):
    """缩放模板与蒙版后做 CCORR_NORMED 匹配。返回 (score, loc, res)。"""
    try:
        if abs(scale - 1.0) < 1e-9:
            tpl, mask = bgr, alpha
        else:
            w = max(8, round(bgr.shape[1] * scale))
            h = max(8, round(bgr.shape[0] * scale))
            if w >= big.shape[1] or h >= big.shape[0]:
                return -1.0, (0, 0), None
            tpl = cv2.resize(bgr, (w, h), interpolation=cv2.INTER_AREA)
            mask = cv2.resize(alpha, (w, h), interpolation=cv2.INTER_AREA)
        res = cv2.matchTemplate(big, tpl, cv2.TM_CCORR_NORMED,
                                mask=mask.astype(np.float32) / 255.0)
        res = np.nan_to_num(res, nan=-1.0, posinf=-1.0, neginf=-1.0).astype(np.float32)
        _, score, _, loc = cv2.minMaxLoc(res)
        return float(score), loc, res
    except cv2.error as e:
        print(f"  匹配失败(scale={scale}): {e}")
        return -1.0, (0, 0), None


def heatmap_png(res):
    r = res.copy()
    r -= r.min()
    if r.max() > 0:
        r /= r.max()
    return cv2.applyColorMap((r * 255).astype(np.uint8), cv2.COLORMAP_JET)


def process_pair(name):
    try:
        big = cv2.imread(str(CAP / f"{name}.jpeg"), cv2.IMREAD_COLOR)
        if big is None:
            print(f"{name[:12]}  大图读取失败")
            return
        bgr, alpha = load_mark(CAP / f"{name}-mark.png")
        if bgr is None:
            print(f"{name[:12]}  mark 读取失败")
            return

        # 单尺度对照
        res0 = cv2.matchTemplate(big, bgr, cv2.TM_CCOEFF_NORMED)
        _, s0, _, l0 = cv2.minMaxLoc(res0)
        s1, l1, _ = match_masked(big, bgr, alpha, 1.0)

        # 多尺度扫描
        scores, locs = [], []
        for sc in SCALES:
            sco, loc, _ = match_masked(big, bgr, alpha, float(sc))
            scores.append(sco)
            locs.append(loc)
        curve = np.array(scores)
        k = int(curve.argmax())
        best_scale = float(SCALES[k])
        best_score, best_loc = float(curve[k]), locs[k]
        masked = curve.copy()
        masked[max(0, k - 1): k + 2] = -1.0
        second = float(masked.max())

        print(f"{name[:12]}  裸1.0={s0:.3f}@{l0}  蒙版1.0={s1:.3f}@{l1}  "
              f"最佳尺度={best_scale:.2f} 分数={best_score:.3f} 次峰={second:.3f} "
              f"峰谷差={float(curve.max() - curve.min()):.3f}")

        # 可视化：匹配框 + 最佳尺度响应热力图
        _, _, res_best = match_masked(big, bgr, alpha, best_scale)
        vis = big.copy()
        w = max(8, round(bgr.shape[1] * best_scale))
        h = max(8, round(bgr.shape[0] * best_scale))
        cv2.rectangle(vis, best_loc, (best_loc[0] + w, best_loc[1] + h), (0, 0, 255), 2)
        cv2.putText(vis, f"s={best_scale:.2f} {best_score:.2f}",
                    (best_loc[0], max(12, best_loc[1] - 4)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 255), 1)
        cv2.imwrite(str(OUT / f"{name[:12]}-match.png"), vis)
        if res_best is not None:
            cv2.imwrite(str(OUT / f"{name[:12]}-heat.png"), heatmap_png(res_best))
    except Exception as e:
        print(f"{name[:12]}  处理失败: {e}")


def main():
    pairs = sorted({p.name.replace("-mark.png", "") for p in CAP.glob("*-mark.png")})
    for name in pairs:
        process_pair(name)
    print(f"\n结果图已写入 {OUT}/")


if __name__ == "__main__":
    main()