import torch
import torch.nn.functional as F


class MatchImageAndMaskSize_Dolphin:
    """Resizes an image and a mask to exactly match a reference image's
    runtime pixel size. Reads reference.shape directly at execution time,
    so it can't drift out of sync the way two separately-wired
    GetImageSize -> width/height links can."""

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image_to_resize": ("IMAGE", {"tooltip": "Image that will be resized to match the reference."}),
                "mask_to_resize": ("MASK", {"tooltip": "Mask that will be resized to match the reference."}),
                "reference_image": ("IMAGE", {"tooltip": "Image whose size image_to_resize and mask_to_resize will be matched to."}),
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK")
    RETURN_NAMES = ("image", "mask")
    FUNCTION = "match"
    CATEGORY = "image/postprocessing"

    def match(self, image_to_resize, mask_to_resize, reference_image):
        target_h = reference_image.shape[1]
        target_w = reference_image.shape[2]

        img = image_to_resize.movedim(-1, 1)
        img = F.interpolate(img, size=(target_h, target_w), mode="bicubic", antialias=True)
        img = img.clamp(0, 1).movedim(1, -1)

        mask = mask_to_resize
        if mask.ndim == 2:
            mask = mask.unsqueeze(0)
        mask = F.interpolate(mask.unsqueeze(1), size=(target_h, target_w), mode="nearest-exact").squeeze(1)

        return (img, mask)


NODE_CLASS_MAPPINGS = {
    "MatchImageAndMaskSize_Dolphin": MatchImageAndMaskSize_Dolphin,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "MatchImageAndMaskSize_Dolphin": "🐬 Match Image & Mask Size (Dolphin)",
}