import sys
from PIL import Image
import numpy as np


def extract_colors(image, max_colors=256):
pixels = np.array(image)

# Reshape the array and get the unique colors
pixels = pixels.reshape(-1, pixels.shape[-1])
unique_colors, counts = np.unique(pixels, axis=0, return_counts=True)

# Sort colors by frequency (most common first)
sorted_colors = unique_colors[np.argsort(-counts)]

# Limit to the max number of colors
selected_colors = sorted_colors[:max_colors]

return selected_colors


def save_palette_image(colors, output_path, width=16, height=16):
image = Image.new("RGBA", (width, height), (0, 0, 0, 0)) # Create a transparent image

# Fill the image with the palette colors
for i, color in enumerate(colors):
if i >= width * height:
break
x = i % width
y = i // width
image.putpixel((x, y), tuple(color))

image.save(output_path)


if __name__ == "__main__":
if len(sys.argv) < 3:
print("Usage: python extract_palette_image.py <input_image_path> <output_palette_image_prefix>")
sys.exit(1)

input_image_path = sys.argv[1]
output_palette_image_prefix = sys.argv[2]

image = Image.open(input_image_path).convert("RGBA")
width, height = image.size

for i in range(0, height, 112):
section = image.crop((0, i, 256, min(i + 112, height)))
colors = extract_colors(section)

# Save first palette image
output_path_1 = f"{output_palette_image_prefix}_{i // 112}_1.png"
save_palette_image(colors[:256], output_path_1)

# Save second palette image if more than 256 colors
if len(colors) > 256:
output_path_2 = f"{output_palette_image_prefix}_{i // 112}_2.png"
save_palette_image(colors[256:], output_path_2)

print("Palettes saved with prefix", output_palette_image_prefix)