class WAS_Image_Select_Color:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "red": ("INT", {"default": 255.0, "min": 0.0, "max": 255.0, "step": 0.1}),
                "green": ("INT", {"default": 255.0, "min": 0.0, "max": 255.0, "step": 0.1}),
                "blue": ("INT", {"default": 255.0, "min": 0.0, "max": 255.0, "step": 0.1}),
                "variance": ("INT", {"default": 10, "min": 0, "max": 255, "step": 1}),
                "use_top_red": ("BOOLEAN", {"default": False}),
                "top_percent": ("FLOAT", {"default": 0.75, "min": 0.01, "max": 0.99, "step": 0.01}),
            },
        }

    RETURN_TYPES = ("IMAGE", "INT", "INT", "INT")
    FUNCTION = "select_color"

    CATEGORY = "WAS Suite/Image/Process"

    def select_color(self, image, red, green, blue, variance, use_top_red, top_percent):

        image, avg_red, avg_green, avg_blue = self.color_pick(tensor2pil(image), red, green, blue, variance, use_top_red, top_percent)

        return (pil2tensor(image), avg_red, avg_green, avg_blue)

    def color_pick(self, image, red, green, blue, variance, use_top_red, top_percent):
        # Convert image to RGB mode
        image = image.convert('RGB')

        # Create a new black image of the same size as the input image
        selected_color = Image.new('RGB', image.size, (0, 0, 0))

        # Get the width and height of the image
        width, height = image.size

        # Initialize variables to store the sum of RGB values and the count of selected pixels
        red_values = []
        green_values = []
        blue_values = []
        selected_pixels = 0

        pix_x = []
        pix_y = []

        # Loop through every pixel in the image
        for x in range(width):
            for y in range(height):
                # Get the color of the pixel
                pixel = image.getpixel((x, y))
                r, g, b = pixel

                # Check if the pixel is within the specified color range
                if ((r >= red-variance) and (r <= red+variance) and
                    (g >= green-variance) and (g <= green+variance) and
                        (b >= blue-variance) and (b <= blue+variance)):

                    # Add the pixel RGB values and positions to the lists
                    pix_x.append(x)
                    pix_y.append(y)
                    red_values.append(r)
                    green_values.append(g)
                    blue_values.append(b)
                    selected_pixels += 1

         # Calculate the average RGB values of the selected pixels
        if selected_pixels==0:
            avg_red = 220
            avg_green = 190
            avg_blue = 190
        else:
            if not use_top_red:
                avg_red = sum(red_values) // selected_pixels
                avg_green = sum(green_values) // selected_pixels
                avg_blue = sum(blue_values) // selected_pixels

                # Set the pixel in the selected_color image to the RGB value of the pixel
                for i in range(selected_pixels):
                    selected_color.putpixel((pix_x[i], pix_y[i]), (red_values[i], green_values[i], blue_values[i]))

            else:
                red_values = np.array(red_values)
                green_values = np.array(green_values)
                blue_values = np.array(blue_values)
                pix_x = np.array(pix_x)
                pix_y = np.array(pix_y)

                # Sort the RGB values based on the red channel value in descending order
                sorted_indices = np.argsort(red_values)[::-1]
                
                # Calculate the number of pixels in the top percent
                top = int(selected_pixels * top_percent)
                
                # Calculate the average RGB values of the top percent pixels
                avg_red = int(np.mean(red_values[sorted_indices[:top+1]]))
                avg_green = int(np.mean(green_values[sorted_indices[:top+1]]))
                avg_blue = int(np.mean(blue_values[sorted_indices[:top+1]]))

                # Set the pixel in the selected_color image to the RGB value of the pixel
                for i in range(top+1):
                    selected_color.putpixel((pix_x[sorted_indices[i]], pix_y[sorted_indices[i]]), (red_values[sorted_indices[i]], green_values[sorted_indices[i]], blue_values[sorted_indices[i]]))
        
        # Return the selected color image
        return selected_color, avg_red, avg_green, avg_blue