# ComfyUI Custom Node: Remove up to 13 images from a batch "at once"
# - Takes a batch (IMAGE)
# - Takes 13 index inputs
# - Removes all specified indices simultaneously (no shifting issues)
# - Index = -1 means "do not remove anything" (ignored)
#
# Save as:
#   ComfyUI/custom_nodes/batch_remove_13_indices/__init__.py
# Restart ComfyUI, then find under: "Batch/Index"

import torch


class BatchRemoveImagesAt13Indices:
    """
    Remove multiple images from a batch at once.

    Notes:
      - Indices are 0-based.
      - Any index == -1 is ignored (meaning: remove nothing for that slot).
      - Out-of-range indices are ignored (with a console warning).
      - Duplicate indices are fine (removed once).
      - If removal would produce an empty batch, this node raises an error.
    """

    @classmethod
    def INPUT_TYPES(cls):
        # 13 separate INT inputs, defaulting to -1
        idx_cfg = {"default": -1, "min": -1, "max": 10**9}
        return {
            "required": {
                "images": ("IMAGE",),

                "remove_index_01": ("INT", idx_cfg),
                "remove_index_02": ("INT", idx_cfg),
                "remove_index_03": ("INT", idx_cfg),
                "remove_index_04": ("INT", idx_cfg),
                "remove_index_05": ("INT", idx_cfg),
                "remove_index_06": ("INT", idx_cfg),
                "remove_index_07": ("INT", idx_cfg),
                "remove_index_08": ("INT", idx_cfg),
                "remove_index_09": ("INT", idx_cfg),
                "remove_index_10": ("INT", idx_cfg),
                "remove_index_11": ("INT", idx_cfg),
                "remove_index_12": ("INT", idx_cfg),
                "remove_index_13": ("INT", idx_cfg),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("images",)
    FUNCTION = "remove"
    CATEGORY = "Batch/Index"

    def remove(
        self,
        images,
        remove_index_01,
        remove_index_02,
        remove_index_03,
        remove_index_04,
        remove_index_05,
        remove_index_06,
        remove_index_07,
        remove_index_08,
        remove_index_09,
        remove_index_10,
        remove_index_11,
        remove_index_12,
        remove_index_13,
    ):
        if not torch.is_tensor(images):
            raise TypeError("Expected 'images' to be a torch Tensor (ComfyUI IMAGE type).")
        if images.ndim != 4:
            raise ValueError(f"Expected 'images' with shape [B,H,W,C], got ndim={images.ndim}.")

        b = int(images.shape[0])
        if b <= 0:
            raise ValueError("Input batch is empty.")

        raw_indices = [
            remove_index_01, remove_index_02, remove_index_03, remove_index_04, remove_index_05,
            remove_index_06, remove_index_07, remove_index_08, remove_index_09, remove_index_10,
            remove_index_11, remove_index_12, remove_index_13,
        ]

        # Build a set of indices to remove (simultaneous removal)
        remove_set = set()
        for idx in raw_indices:
            idx = int(idx)

            # Sentinel: -1 means "do not remove"
            if idx == -1:
                continue

            # Disallow other negative indices (since -1 is reserved)
            if idx < -1:
                print(f"[BatchRemove13] Ignoring invalid negative index {idx} (only -1 is allowed).")
                continue

            # Ignore out-of-range indices rather than clamping (clamping could remove the wrong image)
            if idx < 0 or idx >= b:
                print(f"[BatchRemove13] Ignoring out-of-range index {idx} for batch size {b}.")
                continue

            remove_set.add(idx)

        # If nothing to remove, return original batch unchanged
        if not remove_set:
            return (images,)

        if len(remove_set) >= b:
            raise ValueError(
                f"Removal would produce an empty batch (batch size {b}, requested removals {len(remove_set)})."
            )

        # Keep mask (True = keep, False = remove)
        keep_mask = torch.ones((b,), dtype=torch.bool, device=images.device)
        remove_idx = torch.tensor(sorted(remove_set), dtype=torch.long, device=images.device)
        keep_mask[remove_idx] = False

        out = images[keep_mask]
        return (out,)


NODE_CLASS_MAPPINGS = {
    "BatchRemoveImagesAt13Indices": BatchRemoveImagesAt13Indices,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "BatchRemoveImagesAt13Indices": "Batch: Remove 13 Indices (At Once)",
}