"""Find the source stamp image that best matches a merged clothing image.

This script is intentionally standalone and does not import cmbot project code.
Edit the paths below, or pass them on the command line:

    python match_stamp.py --image output\20260623_094529\TY037\1_TY037.png --stamp-dir D:\stamps

It does not copy, move, rename, delete, or write image files. Results are
printed directly.
"""
from __future__ import print_function

import argparse
import sys
import time
from pathlib import Path

try:
    import cv2
    import numpy as np
    from PIL import Image
except ImportError as exc:
    print("Missing dependency: {}".format(exc))
    print("Install required packages: pip install opencv-python numpy Pillow")
    raise


MERGED_IMAGE = r"D:\chengma\cmbot\output\20260623_094529\TY037\1_TY037.png"
STAMP_DIR = r"D:\chengma\印花和底图\已处理印花\卡通71（66大码200斤 KEKE已上）\横1"

IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
TOP_N = 5

SCALE_MIN = 0.10
SCALE_MAX = 2.00
SCALE_STEP = 0.05
ASPECT_Y_FACTORS = (0.60, 0.70, 0.80, 0.90, 1.00, 1.10, 1.20)

ALPHA_THRESHOLD = 20
LOW_SCORE_WARNING_THRESHOLD = 0.55
TEMPLATE_SCORE_WEIGHT = 0.75
EDGE_SCORE_WEIGHT = 0.25
MIN_MATCH_WIDTH = 24
MIN_MATCH_HEIGHT = 24


def is_image(path):
    return path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS


def iter_images(root):
    root = Path(root)
    if not root.exists():
        return
    for path in sorted(root.rglob("*")):
        if is_image(path):
            yield path


def load_rgba(path):
    """Load with Pillow so Windows Chinese paths are handled reliably."""
    with Image.open(str(path)) as image:
        return np.array(image.convert("RGBA"))


def rgba_to_gray(rgba):
    rgb = rgba[:, :, :3]
    return cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)


def rgba_to_rgb(rgba):
    return rgba[:, :, :3]


def crop_transparent_border(rgba):
    alpha = rgba[:, :, 3]
    ys, xs = np.where(alpha > ALPHA_THRESHOLD)
    if len(xs) == 0 or len(ys) == 0:
        mask = np.ones(alpha.shape, dtype=np.uint8) * 255
        return rgba[:, :, :3], mask

    left = int(xs.min())
    right = int(xs.max()) + 1
    top = int(ys.min())
    bottom = int(ys.max()) + 1
    cropped = rgba[top:bottom, left:right]
    mask = (cropped[:, :, 3] > ALPHA_THRESHOLD).astype(np.uint8) * 255
    return cropped[:, :, :3], mask


def build_scales():
    scales = []
    value = SCALE_MIN
    while value <= SCALE_MAX + 0.0001:
        scales.append(round(value, 4))
        value += SCALE_STEP
    return scales


def build_scale_variants():
    variants = []
    seen = set()
    for scale in build_scales():
        for aspect_y in ASPECT_Y_FACTORS:
            scale_x = scale
            scale_y = round(scale * aspect_y, 4)
            key = (scale_x, scale_y, aspect_y)
            if key in seen:
                continue
            seen.add(key)
            variants.append(
                {
                    "scale_x": scale_x,
                    "scale_y": scale_y,
                    "aspect_y": aspect_y,
                }
            )
    return variants


def safe_match_template(source, template, mask=None):
    if template.shape[0] > source.shape[0] or template.shape[1] > source.shape[1]:
        return None
    if template.shape[0] < MIN_MATCH_HEIGHT or template.shape[1] < MIN_MATCH_WIDTH:
        return None
    if mask is not None and int(np.count_nonzero(mask)) < 9:
        return None

    match_mask = mask
    if match_mask is not None and template.ndim == 3 and match_mask.ndim == 2:
        channels = template.shape[2]
        match_mask = cv2.merge([match_mask] * channels)

    try:
        result = cv2.matchTemplate(source, template, cv2.TM_CCORR_NORMED, mask=match_mask)
    except cv2.error:
        return None

    result = np.asarray(result, dtype=np.float32)
    result[~np.isfinite(result)] = -1.0
    return np.clip(result, 0.0, 1.0)


def resize_for_scale(image, scale_x, scale_y, interpolation):
    height, width = image.shape[:2]
    new_width = max(1, int(round(width * scale_x)))
    new_height = max(1, int(round(height * scale_y)))
    return cv2.resize(image, (new_width, new_height), interpolation=interpolation)


def score_stamp(merged_rgb, merged_edges, stamp_path, scale_variants):
    stamp_rgba = load_rgba(stamp_path)
    stamp_rgb, stamp_mask = crop_transparent_border(stamp_rgba)
    stamp_gray = cv2.cvtColor(stamp_rgb, cv2.COLOR_RGB2GRAY)
    stamp_edges = cv2.Canny(stamp_gray, 80, 160)

    best = None
    for variant in scale_variants:
        scale_x = variant["scale_x"]
        scale_y = variant["scale_y"]
        aspect_y = variant["aspect_y"]
        scaled_rgb = resize_for_scale(stamp_rgb, scale_x, scale_y, cv2.INTER_AREA)
        scaled_gray = resize_for_scale(stamp_gray, scale_x, scale_y, cv2.INTER_AREA)
        scaled_edges = resize_for_scale(stamp_edges, scale_x, scale_y, cv2.INTER_NEAREST)
        scaled_mask = resize_for_scale(stamp_mask, scale_x, scale_y, cv2.INTER_NEAREST)

        template_result = safe_match_template(merged_rgb, scaled_rgb, scaled_mask)
        if template_result is None:
            continue

        edge_result = safe_match_template(merged_edges, scaled_edges, scaled_mask)
        if edge_result is None:
            edge_result = np.zeros(template_result.shape, dtype=np.float32)

        combined = (
            TEMPLATE_SCORE_WEIGHT * template_result
            + EDGE_SCORE_WEIGHT * edge_result
        )
        _min_val, max_val, _min_loc, max_loc = cv2.minMaxLoc(combined)
        x, y = max_loc
        template_score = float(template_result[y, x])
        edge_score = float(edge_result[y, x])
        score = float(max_val)

        if best is None or score > best["score"]:
            best = {
                "score": score,
                "template_score": template_score,
                "edge_score": edge_score,
                "stamp_path": stamp_path,
                "x": int(x),
                "y": int(y),
                "scale_x": float(scale_x),
                "scale_y": float(scale_y),
                "aspect_y": float(aspect_y),
                "width": int(scaled_gray.shape[1]),
                "height": int(scaled_gray.shape[0]),
            }

    return best


def find_best_matches(merged_image, stamp_dir, top_n):
    merged_rgba = load_rgba(merged_image)
    merged_rgb = rgba_to_rgb(merged_rgba)
    merged_gray = rgba_to_gray(merged_rgba)
    merged_edges = cv2.Canny(merged_gray, 80, 160)
    scale_variants = build_scale_variants()

    results = []
    for stamp_path in iter_images(stamp_dir):
        try:
            result = score_stamp(merged_rgb, merged_edges, stamp_path, scale_variants)
        except Exception as exc:
            print("Skipped: {} ({})".format(stamp_path, exc))
            continue
        if result is not None:
            results.append(result)

    results.sort(key=lambda item: item["score"], reverse=True)
    return results[:top_n], len(results)


def print_result(results, matched_count, merged_image, stamp_dir):
    print("Merged image: {}".format(Path(merged_image).resolve()))
    print("Stamp dir: {}".format(Path(stamp_dir).resolve()))
    print("Matched stamp candidates: {}".format(matched_count))
    print("")

    if not results:
        print("No valid image match result.")
        return

    best = results[0]
    print("Best match:")
    print("score: {:.4f}".format(best["score"]))
    print("template_score: {:.4f}".format(best["template_score"]))
    print("edge_score: {:.4f}".format(best["edge_score"]))
    print("stamp: {}".format(Path(best["stamp_path"]).resolve()))
    print("location: x={}, y={}".format(best["x"], best["y"]))
    print("scale_x: {:.4f}".format(best["scale_x"]))
    print("scale_y: {:.4f}".format(best["scale_y"]))
    print("aspect_y: {:.4f}".format(best["aspect_y"]))
    print("size: {}x{}".format(best["width"], best["height"]))
    if best["score"] < LOW_SCORE_WARNING_THRESHOLD:
        print("warning: best score is low; please review manually.")

    print("")
    print("Top {}:".format(len(results)))
    for index, item in enumerate(results, 1):
        line = (
            "{}. {:.4f}  template={:.4f}  edge={:.4f}  "
            "scale_x={:.4f}  scale_y={:.4f}  aspect_y={:.4f}  {}"
        )
        print(
            line.format(
                index,
                item["score"],
                item["template_score"],
                item["edge_score"],
                item["scale_x"],
                item["scale_y"],
                item["aspect_y"],
                Path(item["stamp_path"]).name,
            )
        )


def parse_args(argv):
    parser = argparse.ArgumentParser(
        description="Match one merged clothing image against a stamp directory."
    )
    parser.add_argument("--image", default=MERGED_IMAGE, help="Merged clothing image path")
    parser.add_argument("--stamp-dir", default=STAMP_DIR, help="Stamp image directory")
    parser.add_argument("--top", default=TOP_N, type=int, help="Number of results to print")
    return parser.parse_args(argv)


def main(argv=None):
    args = parse_args(argv or sys.argv[1:])
    merged_image = Path(args.image)
    stamp_dir = Path(args.stamp_dir)

    if not merged_image.is_file():
        print("Merged image not found: {}".format(merged_image))
        return 1
    if not stamp_dir.is_dir():
        print("Stamp directory not found: {}".format(stamp_dir))
        return 1

    started = time.time()
    results, matched_count = find_best_matches(merged_image, stamp_dir, max(1, args.top))
    print_result(results, matched_count, merged_image, stamp_dir)
    print("")
    print("Elapsed: {:.2f}s".format(time.time() - started))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())